Update file model.py

model.py

@@ -33,7 +33,7 @@ class Model:
     def compute_loss(self, input_ids, labels=None, criterion=None):
         lm_logits = self.model(input_ids).logits
 
-        labels = input_ids.to(self.model.device)
+        labels = input_ids.to(GetDevice())
         shift_logits = lm_logits[:, :-1, :].contiguous()
         labels = labels[:, 1:].contiguous()
 
@@ -48,7 +48,7 @@ class Model:
 
         with torch.no_grad():
             encoded_ids = tokenizer.encode(seed_text)
-            input_ids = torch.tensor(encoded_ids).unsqueeze(0).to(self.model.device)
+            input_ids = torch.tensor(encoded_ids).unsqueeze(0).to(GetDevice())
             output = self.model.generate(input_ids, max_length=max_len)
 
             logits = output[0].tolist()