I am trying out a pytorch implementation of Lottery Ticket Hypothesis. For that, I want to freeze the weights in a model that are zero. Is the following a correct way to implement it ?
for name, p in model.named_parameters():
            if 'weight' in name:
                tensor = p.data.cpu().numpy()
                grad_tensor = p.grad.data.cpu().numpy()
                grad_tensor = np.where(tensor == 0, 0, grad_tensor)
                p.grad.data = torch.from_numpy(grad_tensor).to(device)