Fix error for new trainings in LORA

hlinander · Jan 6, 2024 · f982128 · f982128
1 parent 579c812
commit f982128
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/lib/distributed_trainer.py b/lib/distributed_trainer.py
@@ -66,14 +66,15 @@ def distributed_train(requested_configs: List[TrainRun] = None):
         if distributed_train_run is not None:
             try:
                 last_aquired_training = time.time()
-                if (
-                    get_serialization_epoch(
-                        DeserializeConfig(
-                            train_run=distributed_train_run.train_run,
-                            device_id=device_id,
-                        )
+                serialized_epoch = get_serialization_epoch(
+                    DeserializeConfig(
+                        train_run=distributed_train_run.train_run,
+                        device_id=device_id,
                     )
-                    < distributed_train_run.train_run.epochs
+                )
+                if (
+                    serialized_epoch is None
+                    or serialized_epoch < distributed_train_run.train_run.epochs
                 ):
                     do_train_run(distributed_train_run, device_id)
                 else: