Pruning gradients before Optimizer_step
I am working on a project that involves pruning out certain weights in the Neural Network.
I want to know how I can set specific gradient weights to zero before the optimizer performs backpropagation.
I have been able to acheive this in regular Pytorch by using the following code.
loss.backwards()
for name, p in self.named_parameters():
   if 'weight' in name:
        grad_tensor = p.grad.data.cpu()
name = name.split(".")[:-1]
name = ".".join(name)
multiplier = mask[name].cpu()
grad_tensor = grad_tensor * multiplier
p.grad.data = grad_tensor
optimizer.step()
I was told that I need to override the backward_step?