One method is to use a method to mask the gradients after backward() call and before step(). I use that one in my code.

  
      github.com
  
  
    evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L29 16



def initialize_transparent_masks(self):
    def helper(m):
        if isinstance(m, self.DEFAULT_MASKED_MODULES):
            self._mask_dict[m] = torch.ones(m.weight.data.size())
    self.module.apply(helper)


def forward(self, *inputs, **kwargs):
    return self.module.forward(*inputs, **kwargs)


def apply_mask_on_gradients(self):
    def helper(m):
        if isinstance(m, self.DEFAULT_MASKED_MODULES):
            m.weight.grad.data[self._mask_dict[m]!=1]=0
    self.module.apply(helper)


def initiliaze_forward_hooks(self):
    def f_hook(module,inp,out):
        #TODO check inp is a tuple of size 1.
        self._inp_dict[module]=inp[0]
    def helper(m):



  
  
    
    
  
  


A second method could be doing the same with backwardhooks; where you use module.register_backward_hook() to mask the gradients. You can have a flag in your wrapper module like is_masked=True/False and then you can check that flag first thing in your hook: if module.is_masked: do masking else: pass
One tricky thing is to reinitialize your optimizer or gradient history of it after each update on your mask to prevent momentum to come into play.