@jamesproud I tested it out and it looks like you don’t need to mask the gradients. If you mask the weights in every forward pass, then those gradient values will not have an affect on the model learned.
Also it saves computing time since masking the grads takes a bit of time.