Upload MyLLaMa

configure_for_hf.py

@@ -65,12 +65,12 @@ class MyLLaMa(PreTrainedModel):
                 -torch.inf,
             )
             .transpose(0, 1)
-            .cuda()
+            .to(self.model.embed.weight.device)
         )
 
         pad_mask = torch.where(
             tensor == self.model.tokenizer.pad_token_id, False, True
-        ).cuda()
+        ).to(self.model.embed.weight.device)
         outs = namedtuple("output", ["logits", "loss"])
         logits = self.model(tensor, att_mask, pad_mask)["logits"]