fix step count

v2/train.py

@@ -274,7 +274,7 @@ def evaluate_loss(model, rows: np.ndarray, device, batch: int) -> float:
 
 
 def train(model, rows: np.ndarray, val_rows: np.ndarray, args, device) -> list[dict]:
-    steps = args.epochs * len(rows) // args.batch
+    steps = int(args.epochs * len(rows) // args.batch)
     decay = [p for n, p in model.named_parameters() if p.dim() >= 2]
     no_decay = [p for n, p in model.named_parameters() if p.dim() < 2]
     optimizer = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1}, {"params": no_decay, "weight_decay": 0.0}],