Update llm_trainer.py

llm_trainer.py

@@ -7,7 +7,8 @@ class LLMTrainer:
                  model: torch.nn.Module = None,
                  tokenizer: PreTrainedTokenizer | AutoTokenizer = None,
                  model_returns_logits: bool = False):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device_type = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(self.device_type)
 
         if tokenizer is None:
             tokenizer = AutoTokenizer.from_pretrained("gpt2")
@@ -31,7 +32,7 @@ class LLMTrainer:
         with torch.no_grad():
             while generated_tokens.size(1) < length:
 
-                with torch.autocast(device_type=self.device, dtype=torch.bfloat16):
+                with torch.autocast(device_type=self.device_type, dtype=torch.bfloat16):
                     if self.model_returns_logits:
                         logits = self.model(generated_tokens)
                     else: