Update tpu_spawn.py

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -131,9 +131,10 @@ def _move_optimizer_state(self, device: Optional[torch.device] = None) -> None:
         """Moves the state of the optimizers to the TPU if needed."""
         # TODO: `self.root_device` would raise error if called outside the spawn process
         # while training on 8 and more cores.
+        device = device or self.root_device
         for opt in self.optimizers:
             for p, v in opt.state.items():
-                opt.state[p] = apply_to_collection(v, torch.Tensor, move_data_to_device, xm.xla_device())
+                opt.state[p] = apply_to_collection(v, torch.Tensor, move_data_to_device, device)
 
     def _setup_model(self, model: Module) -> Module:
         return model