💩🐛Fixed __iter__ preps

carefree0910 · Oct 2, 2024 · 9e2646d · 9e2646d
1 parent 3e04700
commit 9e2646d
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/core/learn/schema.py b/core/learn/schema.py
@@ -285,14 +285,16 @@ def _iter(self: DataLoader) -> Iterator[tensor_dict_type]:
 
                 return _iter
 
-            base_dataloader = prepared.base_dataloader
-            base_dataloader.data = loader.data
-            base_dataloader.for_inference = loader.for_inference
-            base_dataloader.recover_labels = loader.recover_labels
-            cls_prepared = getattr(prepared.__class__, "_prepared_iter_", False)
+            bdl = prepared.base_dataloader
+            bdl.data = loader.data
+            bdl.for_inference = loader.for_inference
+            bdl.recover_labels = loader.recover_labels
+            cls_prepared = getattr(bdl.__class__, "_prepared_iter_", False)
             if not cls_prepared:
-                prepared.__class__._prepared_iter_ = True
-                prepared.__class__.__iter__ = _iter_factory(prepared.__class__.__iter__)
+                iter_fn = _iter_factory(bdl.__class__.__iter__)
+                bdl.__class__.__iter__ = iter_fn
+                loader.__class__.__iter__ = iter_fn
+                bdl.__class__._prepared_iter_ = True
     return prepared_loaders