Solved by jamesproud in post #5 
              
              
                I’ve included some sample code below. I am trying to mask the gradients of layer1 in this example. As you can see, I’m zeroing out the weights upon initialisation. 
class MaskedLinear(nn.Module):
    # Currently unused. Intended for backward hook.
    def _zero_grad_mask(self, module, grad_input, gr…