KeyError: xxxxxxxxxx when calling optimizer.state_dict()

username: a6m6m
cooked: <p>Hello there,</p>
<p>I am training a deep convolution model unsupervised-ly based on <a href="https://github.com/facebookresearch/deepcluster" rel="nofollow noopener">deepcluster</a> pipeline. The approach alternates as follows:</p>
<p>0- remove classifier module<br>
1- feed inputs<br>
2- compute features<br>
3- cluster inputs based on k means<br>
4- add new classifier module<br>
5- use inputs cluster assignments as pseudo labels to train the model</p>
<p>Since I am adding and removing classifier module on the go, I am using <code>optimizer.add_param_group()</code> to add new optimization parameters for the added classifier module and using <code>del optimizer.param_groups[1]</code> to delete the parameters when the module is removed.</p>
<p>The problem is that when saving the model and the optimizer state I call <code>optimizer.state_dict()</code> to get the optimizer state and save it. However the following error is fired:</p>
<pre><code class="lang-auto">/root/module/class.py in save_model_parameters(self, model_parameters_path, epoch, optimizer)
    254                       'state_dict': self.state_dict()}
    255         if optimizer:
--&gt; 256             model_dict['optimizer'] = optimizer.state_dict()
    257 
    258         torch.save(model_dict, model_parameters_path)

/usr/local/lib/python3.6/dist-packages/torch/optim/optimizer.py in state_dict(self)
     96         # Remap state to use order indices as keys
     97         packed_state = {(param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): v
---&gt; 98                         for k, v in self.state.items()}
     99         return {
    100             'state': packed_state,

/usr/local/lib/python3.6/dist-packages/torch/optim/optimizer.py in &lt;dictcomp&gt;(.0)
     96         # Remap state to use order indices as keys
     97         packed_state = {(param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): v
---&gt; 98                         for k, v in self.state.items()}
     99         return {
    100             'state': packed_state,

KeyError: 140413627951216
</code></pre>
<p>It seems that the problem resides in the following part of <code>optimizer.state_dict()</code> method:<br>
<code>param_mappings[id(k)]</code></p>
<p>The id of a certain parameter in the optimizer.state doesn’t exist in param_mappings.<br>
The error is weird and is raised on some environment (colab, cloud) and not on others (local).<br>
Do you have any thoughs?</p>

username: a6m6m
cooked: <p>I believe that the problem resides in the fact that I am removing param_groups as following:<br>
<code>del optimizer.param_groups[1]</code><br>
This is removing the second param group from the optimizer but it seems that it’s associated parameters in <code>optimizer.state</code> reside leading to an error when matching <code>param_mappings</code> with <code>optimizer.state</code>.<br>
<strong>If that is the case then what is the right way to remove param_groups?</strong> (no documentation regarding this in PyTorch optimizer documentation)</p>

username: Utku_Mert_Topcuoglu
cooked: <p>I also would like to learn the correct way to remove an element from the optimizer.param_groups.</p>