fix: correct epoch computation on resume to avoid redoing whole epoch

astrai/trainer/train_context.py

@@ -101,18 +101,13 @@ class TrainContextBuilder:
                 if checkpoint.config:
                     model_config = checkpoint.config
                 if self._resume:
-                    preloaded_epoch = checkpoint.epoch or cfg.start_epoch
+                    preloaded_epoch = checkpoint.epoch
-                    if checkpoint.consumed_samples > 0:
                     per_step = (
-                            cfg.batch_per_device
+                        cfg.batch_per_device * get_world_size() * cfg.grad_accum_steps
-                            * get_world_size()
-                            * cfg.grad_accum_steps
                     )
                     preloaded_consumed = (
                         checkpoint.consumed_samples // per_step
                     ) * per_step
-                    else:
-                        preloaded_consumed = cfg.start_samples * get_world_size()
                     preloaded_checkpoint = checkpoint
 
         if not model_config and hasattr(cfg.model_fn(), "config"):
@@ -162,6 +157,15 @@ class TrainContextBuilder:
             )
 
         sampler_offset = context.consumed_samples // context.world_size
+
+        if self._resume and sampler_offset > 0:
+            offset = context.world_size - 1
+            num_samples_per_replica = (
+                len(train_dataset) + offset
+            ) // context.world_size
+            if num_samples_per_replica > 0:
+                context.epoch = sampler_offset // num_samples_per_replica
+
         sampler = RDSampler(
             data_source=train_dataset,
             start_epoch=context.epoch,