train agaaa

scripts/train.py

@@ -112,13 +112,17 @@ def main():
         batch["output_router_logits"] = True
         return batch
 
-    # Fixed CustomTrainer class
+    # Fixed CustomTrainer class that handles all possible arguments
     class CustomTrainer(Trainer):
-        def compute_loss(self, model, inputs, return_outputs=False):
-            # Remove num_items_in_batch from inputs if present
-            inputs.pop('num_items_in_batch', None)
+        def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+            # Remove any unexpected arguments
+            inputs = {k: v for k, v in inputs.items() if k not in ['num_items_in_batch']}
             
-            with torch.set_grad_enabled(True):  # Ensure gradients are enabled
+            # Ensure we're in training mode
+            model.train()
+            
+            # Forward pass with gradients
+            with torch.set_grad_enabled(True):
                 outputs = model(**inputs)
                 loss = outputs.loss