How to set certain tensors to zero and save the model?


I would like to manually edit some tensors in the model weights and save the new version as .pth file. Following this discussion (How to manually set the weights in a two layer linear model?), I was able to manually set one of the parameters to zero but I don’t know how it could be saved as a weights file.

checkpoint = './log/best_state.pth'
checkpoint = torch.load(checkpoint)

for name, param in model.named_parameters():
      if name in param_name:
            with torch.no_grad():
                # Set parameter vals to zero
                param[param_idx] = torch.zeros_like(param[param_idx])

      , "xyz.pth")

Using, "xyz.pth") should work, as the manipulated parameter will also be saved to the state_dict.
Do you see any issues with this code? If you want to only store the tensor you could use, '').

1 Like