Yet another fix

SINGROUP · Jan 30, 2024 · 99ad4d1 · 99ad4d1
1 parent 410dc73
commit 99ad4d1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/mlspm/data_loading.py b/mlspm/data_loading.py
@@ -270,7 +270,7 @@ def _collate_batch(batch: Iterable[dict]):
 
     # Switch the batch and channel dimension around
     Xs = list(np.stack(Xs, axis=0).transpose(1, 0, 2, 3, 4))
-    if len(Ys) > 0:
+    if len(Ys) > 0 and len(Ys[0]) > 0:
         Ys = list(np.stack(Ys, axis=0).transpose(1, 0, 2, 3))
     sws = list(np.stack(sws, axis=0).transpose(1, 0, 2, 3))
 

diff --git a/mlspm/utils.py b/mlspm/utils.py
@@ -83,7 +83,7 @@ def _get_distributed() -> Tuple[int, int, int, Optional[dist.ProcessGroup]]:
             world_size = dist.get_world_size()
             local_rank = global_rank = dist.get_rank()
         group = dist.group.WORLD
-    except (RuntimeError, AssertionError):
+    except (RuntimeError, AssertionError, ValueError):
         world_size = 1
         group = None
     if world_size <= 1: