Unfortunately, the nn package has not been built around such approach. So it is quite tricky and I don’t know of any example for this.
You can either do this masking in a differentiable way so that you don’t have to change the gradients. But you are not allowed to change the original parameters inplace. So you won’t be able to use your original approach as you will need another Tensor to store the new values.
You can do this masking in a no_grad block using your current code and set the corresponding gradients to 0 afterwards by hooking like on each parameter trained_model.register_hook(lambda grad: mask * grad).