Correct shape for default grid_sizes

bottler · facebook-github-bot · commit b481cfbd019e · 2021-08-17T05:59:07.000-07:00
Summary: Small fix for omitting this argument.

Reviewed By: nikhilaravi

Differential Revision: D29548610

fbshipit-source-id: f25032fab3faa2f09006f5fcf8628138555f2f20
diff --git a/pytorch3d/ops/points_to_volumes.py b/pytorch3d/ops/points_to_volumes.py
@@ -181,8 +181,11 @@ def add_points_features_to_volume_densities_features(
 
     # init the volumetric grid sizes if uninitialized
     if grid_sizes is None:
-        grid_sizes = torch.LongTensor(list(volume_densities.shape[2:])).to(
-            volume_densities
+        # grid sizes shape (minibatch, 3)
+        grid_sizes = (
+            torch.LongTensor(list(volume_densities.shape[2:]))
+            .to(volume_densities)
+            .expand(volume_densities.shape[0], 3)
         )
 
     # flatten densities and features
