unittest: make test_cross_entropy_loss_2d_out_of_bounds_class_index c…

…ompatible for XPU (#665)

This fixes issues from #653

The case is verifying an expected assertion log raised in kernel. We
have different log in XPU backend, different keyword (we have no
'CUDA'), different index (we are using work item).
Original test case uses CUDA specific calls, and expected assert message
returned from CUDA. The test cases are now rewritten in `test_nn_xpu.py`
to adapt to XPU specific counterparts.

- "test_cross_entropy_loss_2d_out_of_bounds_class_index_xpu_float16",
- "test_cross_entropy_loss_2d_out_of_bounds_class_index_xpu_float32",

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
 )