Hello, first of all thanks for helping.
My code is structured in this way:

I initially create the net structure, with all param_grad=True by default
depending on information on a config file I switch to False the param_grad of all heads but one. For instance the code to switch to False the heat map head is the following:

if config["trainable_heads"]["heatmap"] == False:
      enc_hm = map(lambda x: x[1],
                                  filter(lambda p: p[1].requires_grad and not ("backbone.decoder_delta" in p[0]),
                                         model.named_parameters()))
      for param in enc_hm:
          param.requires_grad = False 

I have 3 other identical snippet for the other 3 heads.
Then I train and save the net.
When I resume, I get the error, which only happens if I try to switch the trainable head. If I do not switch heads, the resume works perfectly.
An attempt I made is the following: before resuming from previous training, I set the layers of the network in the exact same  trainability status as just before saving. I checked it was the case by comparing the name of all trainable parameters just before saving and just before resuming. The 2 list are indeed identical and as I would have expected. Which confuses me a lot.
The error is when I load the optimiser:
            self.model.load_state_dict(checkpoint['state_dict'])
            self.scheduler.load_state_dict(checkpoint['scheduler']) 
            self.optimizer.load_state_dict(checkpoint['optimizer'])

No error is raised in the first 2.