Fix crash when cat-ing empty cuda tensors (#5971)

zou3519 · soumith · commit 9923701a0d74 · 2018-03-23T22:22:39.000-04:00
Fixes #5739. The CUDA path for `torch.cat` was missing a check for the case where all input tensors are empty.
diff --git a/aten/src/THC/generic/THCTensorMath.cu b/aten/src/THC/generic/THCTensorMath.cu
@@ -139,6 +139,11 @@ void THCTensor_(catArray)(THCState *state, THCTensor *result,
     }
   }
 
+  // If all inputs are empty tensors, return an empty tensor
+  if (notEmptyTensor == NULL) {
+    return;
+  }
+
   // In the event that the user specified -1 as the concat dimension, then
   // we want to pick the nDims as dimension to cat along (and thus nDims - 1 as the
   // value due to 0-based indexing). If the nDims is // 0 (i.e. we are catting all
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -1078,6 +1078,9 @@ def test_cat(self):
         z = torch.cat([x, y])
         self.assertEqual(z.size(), (21, SIZE, SIZE))
 
+    def test_cat_empty(self):
+        TestTorch._test_cat_empty(self, use_cuda=True)
+
     def test_bernoulli(self):
         x = torch.tensor([0, 1], dtype=torch.cuda.float32)
         self.assertEqual(x.bernoulli().tolist(), [0, 1])
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -2353,17 +2353,25 @@ def test_cat_scalars(self):
                                      'zero-dimensional.*cannot be concatenated'):
             torch.cat([x, y])
 
-    def test_cat_empty(self):
+    @staticmethod
+    def _test_cat_empty(self, use_cuda=False):
         # FIXME: this is legacy behavior and should be removed
         # when we support empty tensors with arbitrary sizes
-        x = torch.randn(4, 3, 32, 32)
-        empty = torch.randn(0)
+        if use_cuda:
+            dtype = torch.cuda.float32
+        else:
+            dtype = torch.float32
+
+        x = torch.randn((4, 3, 32, 32), dtype=dtype)
+        empty = torch.randn((0,), dtype=dtype)
 
         res1 = torch.cat([x, empty], dim=1)
         res2 = torch.cat([empty, x], dim=1)
         self.assertEqual(res1, res2)
 
-        conv = torch.nn.Conv2d(3, 3, kernel_size=1)
+        conv = torch.nn.Conv2d(3, 3, kernel_size=1).float()
+        if use_cuda:
+            conv = conv.cuda()
         res1 = torch.cat([conv(x), empty], dim=1)
         res2 = torch.cat([empty, conv(x)], dim=1)
         self.assertEqual(res1, res2)
@@ -2375,6 +2383,9 @@ def test_cat_empty(self):
                                      'expected a non-empty list of Tensors'):
             torch.cat([], dim=1)
 
+    def test_cat_empty(self):
+        self._test_cat_empty(self)
+
     def test_stack(self):
         x = torch.rand(2, 3, 4)
         y = torch.rand(2, 3, 4)

Original file line number	Diff line number	Diff line change
`@@ -139,6 +139,11 @@ void THCTensor_(catArray)(THCState state, THCTensor result,`
`139`	`139`	`}`
`140`	`140`	`}`
`141`	`141`
	`142`	`+ // If all inputs are empty tensors, return an empty tensor`
	`143`	`+ if (notEmptyTensor == NULL) {`
	`144`	`+ return;`
	`145`	`+ }`
	`146`	`+`
`142`	`147`	`// In the event that the user specified -1 as the concat dimension, then`
`143`	`148`	`// we want to pick the nDims as dimension to cat along (and thus nDims - 1 as the`
`144`	`149`	`// value due to 0-based indexing). If the nDims is // 0 (i.e. we are catting all`