[fix] torch.multinomial : fix for 0 size dim (#43775)

Summary: Fixes pytorch/pytorch#43768 TO-DO: * [x] Add test Pull Request resolved: pytorch/pytorch#43775 Reviewed By: ZolotukhinM Differential Revision: D23421979 Pulled By: ngimel fbshipit-source-id: 949fcdd30f18d17ae1c372fa6ca6a0b8d0d538ce
zhulu506 · Aug 31, 2020 · 0394c5a · 0394c5a
1 parent 3c8b1d7
commit 0394c5a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/aten/src/ATen/native/Distributions.cpp b/aten/src/ATen/native/Distributions.cpp
@@ -454,6 +454,7 @@ Tensor& multinomial_out(Tensor& result, const Tensor& self, int64_t n_sample, bo
   if (self.dim() > 1) {
     int64_t n_dist = self.size(-2);
     result.resize_({n_dist, n_sample});
+    if (n_dist == 0) { return result; };
   } else {
     result.resize_({n_sample});
   }

diff --git a/aten/src/ATen/native/cuda/MultinomialKernel.cu b/aten/src/ATen/native/cuda/MultinomialKernel.cu
@@ -395,6 +395,8 @@ void multinomial_kernel_impl(Tensor& result, const Tensor& self, const int64_t n
     }
   });
 
+  AT_CUDA_CHECK(cudaGetLastError());
+
   if (inputSize == 1) {
     result.resize_({n_sample});
   }

diff --git a/test/test_torch.py b/test/test_torch.py
@@ -17942,13 +17942,19 @@ def test(probs, replacement):
             test(y, True)
             test(z, True)
 
-    def test_multinomial_empty(self, device):
-        probs = torch.ones(0, 3)
-        num_samples = 1
+    def _test_multinomial_empty(self, device, replacement, num_samples):
+        probs = torch.ones(0, 3, device=device)
         expected = torch.empty(0, num_samples, dtype=torch.int64)
-        for replacement in (True, False):
-            out = torch.multinomial(probs, num_samples=num_samples, replacement=replacement)
-            self.assertEqual(out, expected)
+        out = torch.multinomial(probs, num_samples=num_samples, replacement=replacement)
+        self.assertEqual(out, expected)
+
+    def test_multinomial_empty_w_replacement(self, device):
+        self._test_multinomial_empty(device, True, 1)
+        self._test_multinomial_empty(device, True, 2)
+
+    def test_multinomial_empty_wo_replacement(self, device):
+        self._test_multinomial_empty(device, False, 1)
+        self._test_multinomial_empty(device, False, 2)
 
     def _generate_input(self, shape, dtype, device, with_extremal):
         if shape == ():
-Original file line number
+Diff line change
@@ Expand Up @@
         }
       });
+      AT_CUDA_CHECK(cudaGetLastError());
       if (inputSize == 1) {
         result.resize_({n_sample});
       }
@@ Expand Down @@