Update tpu_spawn.py

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -132,12 +132,9 @@ 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.
-        if device:
-            raise ValueError(f"device should be None" f" found: {device}.")
-        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, device)
+                opt.state[p] = apply_to_collection(v, torch.Tensor, move_data_to_device, device or self.root_device)
 
     def _setup_model(self, model: Module) -> Module:
         return model