Error using octree with multiscale_type='cat'

[patmjen]

First, thank you for the great library!

Using OctreeGrid with multiscale_type='cat' gave me an error when calling interpolate at L273

kaolin-wisp/wisp/models/grids/octree_grid.py

Line 273 in 7ca5c6f

return feats.reshape(batch, num_samples, self.feature_dim)

The problem is that self.feature_dim is too small as the entries have self.feature_dim * num_feats dimensions.

I changed the code as follows, which seemed to fix it.

diff --git a/wisp/models/grids/octree_grid.py b/wisp/models/grids/octree_grid.py
index 7c6a8bc..95ec5e3 100644
--- a/wisp/models/grids/octree_grid.py
+++ b/wisp/models/grids/octree_grid.py
@@ -263,14 +263,17 @@ class OctreeGrid(BLASGrid):

             if self.multiscale_type == 'sum':
                 feats = feats.reshape(*feats.shape[:-1], num_feats, self.feature_dim)
+                out_features = self.feature_dim
                 if self.training:
                     feats = feats.sum(-2)
                 else:
                     feats = feats.sum(-2)
+            else:  # self.multiscale_type == 'cat'
+                out_features = self.feature_dim * num_feats

             timer.check("aggregate")

-            return feats.reshape(batch, num_samples, self.feature_dim)
+            return feats.reshape(batch, num_samples, out_features)

     def raymarch(self, rays, level=None, num_samples=64, raymarch_type='voxel'):
         """Mostly a wrapper over OctreeAS.raymarch. See corresponding function for more details.

[orperel]

Thanks a lot for reporting this one @patmjen ! :)