Save Custom Parameters in Module using state_dict()

Currently when I use

torch.save(lstm_net.state_dict())

it does not save custom parameters I have added to the module. It only saves the weights and biases of the lstm and dense layers.

How do I fix it so it does?

are your custom parameters registered in the constructor with self.param_name = nn.Parameter(...)?
The key is to wrap them in nn.Parameter, or else they wont be under the purview of state_dict or model.parameters()

2 Likes

That was it. Thank you