bug fixes

src/model.py

@@ -359,7 +359,8 @@ class ChessForCausalLM(PreTrainedModel):
             shift_labels = labels[..., 1:].contiguous()
             
             # Flatten for cross-entropy
-            loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            # loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
             loss = loss_fct(
                 shift_logits.view(-1, shift_logits.size(-1)),
                 shift_labels.view(-1),

src/train.py

@@ -35,10 +35,6 @@ def parse_args():
     )
     
     # Model arguments
-    parser.add_argument(
-        "--vocab_size", type=int, default=1200,
-        help="Vocabulary size"
-    )
     parser.add_argument(
         "--n_embd", type=int, default=128,
         help="Embedding dimension"