Fix

SINGROUP · Jan 30, 2024 · 9cde4d8 · 9cde4d8
1 parent 7d78902
commit 9cde4d8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mlspm/data_loading.py b/mlspm/data_loading.py
@@ -208,7 +208,7 @@ def _rotate_and_stack(src: Iterable[dict], reverse: bool = False) -> Generator[d
         for key in sample.keys():
             if key[-3:] in ["jpg", "png"]:
                 num = key[:-4]
-                if '.' in key:
+                if '.' in num:
                     iz, channel = [int(n) for n in num.split('.')]
                 else:
                     iz, channel = int(num), 0