Update litgpt/model.py

litgpt/model.py

@@ -16,10 +16,12 @@ from litgpt.config import Config
 
 
 class GPT(nn.Module):
-    def __init__(self, config: Config) -> None:
+    def __init__(self, config: Config, device=None) -> None:
         super().__init__()
         assert config.padded_vocab_size is not None
         self.config = config
+        if device is not None:
+            self.device = device
         if self.config.asr_adapter == "mlp":
             print("Using MLP adapter for ASR feature")
             self.whisper_adapter = nn.Linear(config.whisper_adapter_dim, config.n_embd)