Update difpoint/src/modules/dense_motion.py

difpoint/src/modules/dense_motion.py

@@ -30,7 +30,7 @@ class DenseMotionNetwork(nn.Module):
         bs, _, d, h, w = feature.shape  # (bs, 4, 16, 64, 64)
         identity_grid = make_coordinate_grid((d, h, w), ref=kp_source)  # (16, 64, 64, 3)
         identity_grid = identity_grid.view(1, 1, d, h, w, 3)  # (1, 1, d=16, h=64, w=64, 3)
-        coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3)
+        coordinate_grid = identity_grid - kp_driving.reshape(bs, self.num_kp, 1, 1, 1, 3)
 
         k = coordinate_grid.shape[1]