I am experimenting with randomly removing individual weights after the model has been trained. In order to restore the original state of the network after these experiments, I cache them before. However, the restoring does not work as expected, can somebody point out my mistake?
[...]
cache = self.state_dict()
for layer in self.layers.parameters():
    print(layer.data[index])  # tensor(0.0034)

    layer.data[index] = 0
    print(layer.data[index])  # tensor(0.)

    [...] #evaluate the model using the disabled weight

    self.load_state_dict(cache)
    print(layer.data[index])  # tensor(0.)