fix: empty depth stack

mosure · Oct 31, 2024 · 67e2938 · 67e2938
1 parent eb24925
commit 67e2938
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ffi/python/bevy_zeroverse_dataloader/__init__.py b/ffi/python/bevy_zeroverse_dataloader/__init__.py
@@ -366,7 +366,11 @@ def __getitem__(self, idx):
                 depth_tensors.append(depth_tensor)
 
         color_tensor = torch.stack(color_tensors, dim=0)
-        depth_tensor = torch.stack(depth_tensors, dim=0)
+
+        if len(depth_tensors) == 0:
+            depth_tensor = torch.zeros_like(color_tensor)
+        else:
+            depth_tensor = torch.stack(depth_tensors, dim=0)
 
         meta_tensors['color'] = color_tensor
         meta_tensors['depth'] = depth_tensor