adapt test_cross_entropy_loss_2d_out_of_bounds_class_index to xpu

test/xpu/run_test_with_skip.py

@@ -1275,9 +1275,6 @@ def launch_test(test_case, skip_list=None, exe_list=None):
     # NotImplementedError: Could not run 'aten::_indices' with arguments from the 'SparseXPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build).
     "test_EmbeddingBag_sparse_cuda",
     "test_Embedding_sparse_cuda",
-    # AssertionError: 'XPU error: device-side assert triggered' not found in '  File "<string>", line 8\n    def test_cross_entropy_loss_2d_out_of_bounds_class_index(self):\n    ^\nIndentationError: expected an indented block\n'
-    "test_cross_entropy_loss_2d_out_of_bounds_class_index_xpu_float16",
-    "test_cross_entropy_loss_2d_out_of_bounds_class_index_xpu_float32",
     # AssertionError: MultiheadAttention does not support NestedTensor outside of its fast path. The fast path was not hit because some Tensor argument's device is neither one of cpu, cuda or privateuseone
     "test_TransformerEncoderLayer_empty_xpu",
     "test_transformerencoderlayer_xpu_float16",

test/xpu/test_nn_xpu.py

@@ -2240,6 +2240,45 @@ def upsamplingBiMode2d(self, device, antialias, align_corners, mode, memory_form
                 self.assertEqual(a_xpu.grad, a_cpu.grad)
 TestNNDeviceType.test_upsamplingBiMode2d = upsamplingBiMode2d
 
+@dtypes(torch.float16, torch.float32)
+def _test_cross_entropy_loss_2d_out_of_bounds_class_index(self, device, dtype):
+    from torch.testing._internal.common_utils import TestCase
+    # Test for issue #117532
+    # Run in a different process to prevent the device-side assert from affecting other tests
+    stderr = TestCase.runWithPytorchAPIUsageStderr(f"""\
+#!/usr/bin/env python3
+
+import torch
+import torch.nn.functional as F
+from torch.testing._internal.common_utils import (run_tests, TestCase)
+
+class TestThatContainsCUDAAssert(TestCase):
+    def test_cross_entropy_loss_2d_out_of_bounds_class_index(self):
+        device = '{str(device)}'
+        dtype = {str(dtype).strip("'")}
+        ignore_index = 255
+        b = 10
+        n_classes = 3
+        w = 768
+        h = 1024
+        pred = torch.randn(b, n_classes, w, h, dtype=dtype, device=device)
+        labels = torch.zeros(b, w, h, dtype=torch.int64, device=device)
+        labels[5, 200, 200] = ignore_index
+        # Set invalid class index
+        labels[5, 200, 200] = 254
+
+        x = F.cross_entropy(
+            pred, labels, reduction="none", ignore_index=ignore_index
+        )
+        torch.xpu.synchronize()
+
+
+if __name__ == '__main__':
+    run_tests()
+        """)
+    self.assertIn('Assertion `cur_target >= 0 && cur_target < n_classes` failed', stderr)
+TestNNDeviceType.test_cross_entropy_loss_2d_out_of_bounds_class_index = _test_cross_entropy_loss_2d_out_of_bounds_class_index
+
 instantiate_device_type_tests(
     TestNNDeviceType, globals(), only_for="xpu", allow_xpu=True
 )