diff --git a/tfmplayground/training/train.py b/tfmplayground/training/train.py
index e428b0e..2596bc4 100644
--- a/tfmplayground/training/train.py
+++ b/tfmplayground/training/train.py
@@ -32,8 +32,7 @@ def train(
     """
     trains model on prior batches for given epochs
 
-    steps with non-finite values are skipped, and an epoch where every step is
-    skipped raises
+    batches with non-finite values are retried, and 100 consecutive skips raise
 
     Parameters
     ----------
@@ -48,7 +47,7 @@ def train(
     batch_size : int
         number of tables per batch
     steps_per_epoch : int
-        number of batches that make up one epoch
+        number of optimizer updates per epoch
     lr : float
         learning rate for optimizer
     grad_clip : float
@@ -65,6 +64,8 @@ def train(
     TabularFoundationModel
         model after last epoch
     """
+    if steps_per_epoch < 1:
+        raise ValueError("steps_per_epoch must be at least 1")
     if callbacks is None:
         callbacks = []
     if not device:
@@ -83,7 +84,11 @@ def train(
             optimizer.train()
             total_loss = 0.0
             num_valid = 0
-            for _ in range(steps_per_epoch):
+            consecutive_skips = 0
+            while num_valid < steps_per_epoch:
+                if consecutive_skips >= 100:
+                    raise RuntimeError(f"100 consecutive batches in epoch {epoch} had non-finite values")
+                consecutive_skips += 1
                 x_train, y_train, x_test, y_test = next(batches)
                 x_train = x_train.to(device)
                 y_train = y_train.to(device)
@@ -117,15 +122,12 @@ def train(
 
                 loss = losses.mean()
                 loss.backward()
-                total_loss += loss.cpu().detach().item()
-                num_valid += 1
-
                 torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                 optimizer.step()
                 optimizer.zero_grad()
-
-            if num_valid == 0:
-                raise RuntimeError(f"all {steps_per_epoch} steps in epoch {epoch} had non-finite values")
+                total_loss += loss.cpu().detach().item()
+                num_valid += 1
+                consecutive_skips = 0
 
             end_time = time.time()
             mean_loss = total_loss / num_valid