.dense fix cuda cpu mismatch (#319)

NVIDIA · May 11, 2021 · a445cf6 · a445cf6
1 parent c4470e2
commit a445cf6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/MinkowskiEngine/MinkowskiSparseTensor.py b/MinkowskiEngine/MinkowskiSparseTensor.py
@@ -492,7 +492,7 @@ def dense(self, shape=None, min_coordinate=None, contract_stride=True):
                 shape = torch.Size([shape[0], self._F.size(1), *[s for s in shape[2:]]])
 
         # Use int tensor for all operations
-        tensor_stride = torch.IntTensor(self.tensor_stride)
+        tensor_stride = torch.IntTensor(self.tensor_stride).to(self.device)
 
         # New coordinates
         batch_indices = self.C[:, 0]
@@ -527,7 +527,7 @@ def dense(self, shape=None, min_coordinate=None, contract_stride=True):
         nchannels = self.F.size(1)
         if shape is None:
             size = coords.max(0)[0] + 1
-            shape = torch.Size([batch_indices.max() + 1, nchannels, *size.numpy()])
+            shape = torch.Size([batch_indices.max() + 1, nchannels, *size.cpu().numpy()])
 
         dense_F = torch.zeros(shape, dtype=self.F.dtype, device=self.F.device)