Add method to TPUSpawn plugin to override how models are setup (Lightning-AI#10039)

Author: kaushikb11

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -20,6 +20,7 @@
 
 import torch
 import torch.multiprocessing as mp
+from torch.nn import Module
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
@@ -118,6 +119,9 @@ def pre_dispatch(self):
     def setup(self) -> None:
         self.create_mp_queue()
 
+    def _setup_model(self, model: Module) -> Module:
+        return model
+
     def create_mp_queue(self):
         self.start_method = "fork"
         smp = mp.get_context(self.start_method)