fixed zero numel init bug

yanbing-j · Apr 16, 2020 · 1eef2be · 1eef2be
1 parent 66105d4
commit 1eef2be
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 8 deletions.
diff --git a/csrc/cpu/scatter_cpu.cpp b/csrc/cpu/scatter_cpu.cpp
@@ -43,8 +43,11 @@ scatter_cpu(torch::Tensor src, torch::Tensor index, int64_t dim,
     arg_out_data = arg_out.value().data_ptr<int64_t>();
   }
 
-  if (index.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return std::make_tuple(out, arg_out);
+  }
 
   auto B = 1;
   for (auto i = 0; i < dim; i++)

diff --git a/csrc/cpu/segment_coo_cpu.cpp b/csrc/cpu/segment_coo_cpu.cpp
@@ -52,8 +52,11 @@ segment_coo_cpu(torch::Tensor src, torch::Tensor index,
     arg_out = torch::zeros(sizes, out.options());
   }
 
-  if (index.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return std::make_tuple(out, arg_out);
+  }
 
   auto B = index.numel() / src.size(dim);
   auto E = src.size(dim);
@@ -158,8 +161,11 @@ torch::Tensor gather_coo_cpu(torch::Tensor src, torch::Tensor index,
     out = torch::empty(sizes, src.options());
   }
 
-  if (index.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return out;
+  }
 
   auto B = index.numel() / out.size(dim);
   auto E = index.size(dim);

diff --git a/csrc/cpu/segment_csr_cpu.cpp b/csrc/cpu/segment_csr_cpu.cpp
@@ -44,8 +44,11 @@ segment_csr_cpu(torch::Tensor src, torch::Tensor indptr,
     arg_out_data = arg_out.value().data_ptr<int64_t>();
   }
 
-  if (src.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return std::make_tuple(out, arg_out);
+  }
 
   auto N = out.size(dim) * (indptr.numel() / indptr.size(-1));
   auto K = out.numel() / N;
@@ -120,8 +123,11 @@ torch::Tensor gather_csr_cpu(torch::Tensor src, torch::Tensor indptr,
     out = torch::empty(sizes, src.options());
   }
 
-  if (src.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return out;
+  }
 
   auto N = src.size(dim) * (indptr.numel() / indptr.size(-1));
   auto K = src.numel() / N;

diff --git a/csrc/cuda/scatter_cuda.cu b/csrc/cuda/scatter_cuda.cu
@@ -99,8 +99,11 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,
     arg_out_data = arg_out.value().data_ptr<int64_t>();
   }
 
-  if (index.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return std::make_tuple(out, arg_out);
+  }
 
   auto B = 1;
   for (auto i = 0; i < dim; i++)

diff --git a/csrc/cuda/segment_csr_cuda.cu b/csrc/cuda/segment_csr_cuda.cu
@@ -135,8 +135,11 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
     arg_out_data = arg_out.value().data_ptr<int64_t>();
   }
 
-  if (src.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return std::make_tuple(out, arg_out);
+  }
 
   auto N = out.size(dim) * (indptr.numel() / indptr.size(-1));
   auto K = out.numel() / N;
@@ -251,8 +254,11 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr,
     out = torch::empty(sizes, src.options());
   }
 
-  if (src.numel() == 0)
+  if (src.numel() == 0) {
+    if (!optional_out.has_value())
+      out.fill_(0);
     return out;
+  }
 
   auto N = src.size(dim) * (indptr.numel() / indptr.size(-1));
   auto K = src.numel() / N;