add asserts

mariogeiger · Aug 8, 2019 · c91bd59 · c91bd59
1 parent 284f27c
commit c91bd59
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/se3cnn/point/operations.py b/se3cnn/point/operations.py
@@ -14,6 +14,7 @@ def forward(self, features, geometry):
         :param geometry: tensor [batch, point, xyz]
         :return:         tensor [batch, point, channel]
         """
+        assert features.size()[:2] == geometry.size()[:2], "features size ({}) and geometry size ({}) should match".format(features.size(), geometry.size())
         rb = geometry.unsqueeze(1)  # [batch, 1, b, xyz]
         ra = geometry.unsqueeze(2)  # [batch, a, 1, xyz]
         k = self.kernel(rb - ra)  # [batch, a, b, i, j]
@@ -31,6 +32,7 @@ def forward(self, features, geometry):
         :param geometry: tensor [batch, point, xyz]
         :return:         tensor [batch, point, point, channel]
         """
+        assert features.size()[:2] == geometry.size()[:2], "features size ({}) and geometry size ({}) should match".format(features.size(), geometry.size())
         rb = geometry.unsqueeze(1)  # [batch, 1, b, xyz]
         ra = geometry.unsqueeze(2)  # [batch, a, 1, xyz]
         k = self.kernel(rb - ra)  # [batch, a, b, i, j]
@@ -50,6 +52,7 @@ def forward(self, features, geometry):
         :param geometry: tensor [batch, point, xyz]
         :return:         tensor [batch, point, channel]
         """
+        assert features.size()[:2] == geometry.size()[:2], "features size ({}) and geometry size ({}) should match".format(features.size(), geometry.size())
         batch, n, _ = geometry.size()
 
         rb = geometry.unsqueeze(1)  # [batch, 1, b, xyz]