Updating the 1D-1D shape check for NLL loss.

mansiag05 · pytorchmergebot · commit 1398d84b4226 · 2025-09-05T18:14:17.000Z
diff --git a/aten/src/ATen/native/LossNLL.cpp b/aten/src/ATen/native/LossNLL.cpp
@@ -47,20 +47,20 @@ TORCH_META_FUNC(nll_loss_forward)
   TORCH_CHECK(
       target.dim() <= 1,
       "0D or 1D target tensor expected, multi-target not supported");
-
-  auto no_batch_dim = self.dim() == 1  && target.dim() == 0;
+  if (self.dim() == 1 && target.dim() == 1) {
+      TORCH_CHECK_VALUE(
+          target.size(0) == 1,
+          "For 1D input, 1D target must have size 1, but got target size: ",
+          target.size(0));
+  }
   TORCH_CHECK(
-      no_batch_dim || (self.size(0) == target.size(0)),
+      self.dim() == 1 || (self.size(0) == target.size(0)),
       "size mismatch (got input: ",
       self.sizes(),
       ", target: ",
       target.sizes(),
       ")")
 
-  TORCH_CHECK_VALUE(
-        !(self.dim() == 1 && target.dim() == 1),
-        "Expected target to be 0D when input is 1D, but got target.dim =", target.dim());
-
   const auto n_classes = self.size(-1);
 
   TORCH_CHECK(
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -12048,11 +12048,10 @@ def test_softmax_bfloat16(self, device):
             # test softmax with large input value which causes exp() to overflow
             _test_bfloat16_ops(self, torch.nn.Softmax(dim=dim), device, inp_dims=(16, 33, 15, 16), prec=0.05, scale_factor=1000.0)
 
-    def test_nll_loss_incompatible_1d_input_1d_target(self, device):
-        x = torch.randn((10,), device=device)  # 1D input: single sample, 10 classes
-        t = torch.zeros((10,), dtype=torch.int64, device=device)  # 1D target: 10 class indices
-
-        with self.assertRaisesRegex(ValueError, "Expected target to be 0D when input is 1D"):
+    def test_nll_loss_1d_input_1d_target_invalid_size(self, device):
+        x = torch.randn(10, device=device)
+        t = torch.randint(0, 10, (3,), dtype=torch.int64, device=device)
+        with self.assertRaisesRegex(ValueError, "For 1D input, 1D target must have size 1"):
             F.nll_loss(x, t)
 
     def test_nll_loss_mismatched_batch(self, device):
