return if input numel is 0

pytorch · Jan 7, 2021 · 76aa2ab · 76aa2ab
1 parent 4ef5707
commit 76aa2ab
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/aten/src/ATen/native/LossMultiMargin.cpp b/aten/src/ATen/native/LossMultiMargin.cpp
@@ -212,6 +212,10 @@ void multi_margin_loss_backward_out_cpu_template(
   grad_input.resize_as_(input);
   TORCH_CHECK(grad_input.is_contiguous(), "grad_input must be contiguous");
 
+  if (input.numel() == 0) {
+    return;
+  }
+
   auto input_contiguous = input.contiguous();
   auto target_contiguous = target.contiguous();
   auto weight_contiguous = weight.contiguous();