bug fix: restore_optimizers correctly handles non-mapping values in optimizer.state.values() (#11757)

Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com>

Author: circlecrystal

CHANGELOG.md

@@ -573,6 +573,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - The `RichProgressBar` now correctly shows the `on_epoch` logged values on train epoch end ([#11689](https://github.com/PyTorchLightning/pytorch-lightning/pull/11689))
 
 
+- Fixed `restore_optimizers` for mapping states ([#11757](https://github.com/PyTorchLightning/pytorch-lightning/pull/11757))
+
+
 - Fixed check for available modules ([#11526](https://github.com/PyTorchLightning/pytorch-lightning/pull/11526))
 
 

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -282,10 +282,13 @@ def restore_optimizers(self) -> None:
             # move optimizer to GPU 1 weight at a time
             # avoids OOM
             if self.trainer.root_gpu is not None:
-                for state in optimizer.state.values():
-                    for k, v in state.items():
-                        if isinstance(v, torch.Tensor):
-                            state[k] = v.cuda(self.trainer.root_gpu)
+                for param, state in optimizer.state.items():
+                    if isinstance(state, dict):
+                        for k, v in state.items():
+                            if isinstance(v, torch.Tensor):
+                                state[k] = v.cuda(self.trainer.root_gpu)
+                    elif isinstance(state, torch.Tensor):
+                        optimizer.state[param] = state.cuda(self.trainer.root_gpu)
 
     def restore_lr_schedulers(self) -> None:
         """Restores the learning rate scheduler states from the pre-loaded checkpoint."""