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

1

Hello,

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])

                torch.save(model.state_dict(), "xyz.pth")
2

Using torch.save(model.state_dict(), "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 torch.save(model.layer.param, 'path.pt').

1 Like