fix empty

zasdfgbnm · zasdfgbnm · commit ae7760cf50bb · 2021-05-24T10:15:34.000-07:00
diff --git a/aten/src/ATen/native/cuda/IndexKernel.cu b/aten/src/ATen/native/cuda/IndexKernel.cu
@@ -373,6 +373,10 @@ template <typename mask_t>
 void masked_scatter_cuda_impl(Tensor& self, const Tensor& mask, const Tensor& source){
   auto srcSize = source.numel();
 
+  if (srcSize == 0) {
+    return;
+  }
+
   auto mask_cont = mask.contiguous();
 
   // Use a prefix sum to determine the output locations of the masked elements
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -5708,11 +5708,21 @@ def test_masked_scatter(self, device, dtype):
                     with self.assertRaises(RuntimeError):
                         dest.masked_scatter_(mask, src)
 
+                # empty tensor
+                dest = torch.empty((5, 0, 5), dtype=dt, device=device)
+                mask = torch.empty_like(dest, dtype=maskType, device=device)
+                src = torch.empty((0,), dtype=dt, device=device)
+                dest.masked_scatter_(mask, src)
+
+                dest = torch.empty((5, 0, 5), dtype=dt, device=device)
+                mask = torch.empty((5, 1, 5), dtype=maskType, device=device)
+                src = torch.empty((0,), dtype=dt, device=device)
+                dest.masked_scatter_(mask, src)
 
         if self.device_type != 'cuda':
-            self.assertEqual(len(w), 3)
+            self.assertEqual(len(w), 5)
         else:
-            self.assertEqual(len(w), 2)
+            self.assertEqual(len(w), 4)
 
         warn = 'masked_scatter_ received a mask with dtype torch.uint8,'
         for wi in w:
