`NormConstraint` should be instantiated with a `max_norm`, rather than it being passed to `constrain_param`. And why are you not using `constrain_update` when you have that method and you are actually constraining updates?