train aga

scripts/train.py

@@ -137,8 +137,12 @@ def main():
     
     # Test forward/backward pass before training
     print("Testing gradient flow...")
-    test_batch = next(iter(DataLoader(tokenized_dataset, batch_size=1)))
-    test_batch = {k: v.to(model.device) for k, v in test_batch.items()}
+    test_loader = DataLoader(tokenized_dataset, batch_size=1, collate_fn=data_collator)
+    test_batch = next(iter(test_loader))
+    
+    # Move batch to model's device
+    device = next(model.parameters()).device
+    test_batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in test_batch.items()}
     
     model.train()
     outputs = model(**test_batch)