Use epoch loop loader for transfer to correctly handle end_of_epoch (#2965)

pplantinga · web-flow · commit 308527c03d9f · 2025-08-13T10:46:56.000+02:00
diff --git a/speechbrain/utils/epoch_loop.py b/speechbrain/utils/epoch_loop.py
@@ -12,6 +12,7 @@
 from .checkpoints import (
     mark_as_loader,
     mark_as_saver,
+    mark_as_transfer,
     register_checkpoint_hooks,
 )
 
@@ -65,6 +66,7 @@ def _save(self, path):
             fo.write(str(self.current))
 
     @mark_as_loader
+    @mark_as_transfer
     def _recover(self, path, end_of_epoch=True):
         # NOTE: end_of_epoch = True by default so that when
         #  loaded in parameter transfer, this starts a new epoch.
@@ -183,6 +185,7 @@ def _save(self, path):
             )
 
     @mark_as_loader
+    @mark_as_transfer
     def _recover(self, path, end_of_epoch=True, device=None):
         del device  # Not used.
         with open(path, encoding="utf-8") as fi:
