Upload MyLLaMa

configure_for_hf.py

@@ -64,13 +64,12 @@ class MyLLaMa(PreTrainedModel):
                 -torch.inf,
             )
             .transpose(0, 1)
-            .cuda(),
-        )
-        pad_mask = (
-            torch.where(
-                tensor == self.model.tokenizer.pad_token_id, False, True
-            ).cuda(),
+            .cuda()
         )
+
+        pad_mask = torch.where(
+            tensor == self.model.tokenizer.pad_token_id, False, True
+        ).cuda()
         logits = self.model(tensor, att_mask, pad_mask)["logits"]
         if labels is not None:
             loss = nn.functional.cross_entropy(logits, labels)