refine the xpu skip func

Author: zxd1997066

torchao/testing/utils.py

@@ -98,22 +98,24 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
-def skip_if_no_xpu(message=None):
-    """Decorator to skip tests on ROCm platform with custom message.
+def skip_if_no_xpu():
+    try:
+        import pytest
 
-    Args:
-        message (str, optional): Additional information about why the test is skipped.
-    """
-    import unittest
+        has_pytest = True
+    except ImportError:
+        has_pytest = False
+        import unittest
 
     def decorator(func):
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             if not torch.xpu.is_available():
-                skip_message = "Skipping the test in XPU"
-                if message:
-                    skip_message += f": {message}"
-                unittest.skip(skip_message)
+                skip_message = "No XPU available"
+                if has_pytest:
+                    pytest.skip(skip_message)
+                else:
+                    unittest.skip(skip_message)
             return func(*args, **kwargs)
 
         return wrapper
@@ -123,19 +125,39 @@ def wrapper(*args, **kwargs):
 
 def skip_if_xpu(message=None):
     """
-    Decorator to skip tests if XPU is available.
+    Decorator to skip tests on XPU platform with custom message.
 
     Args:
         message (str, optional): Additional information about why the test is skipped.
     """
+    try:
+        import pytest
+
+        has_pytest = True
+    except ImportError:
+        has_pytest = False
+        import unittest
 
     def decorator(func):
-        reason = "Skipping the test on XPU"
-        if message:
-            reason += f": {message}"
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if torch.xpu.is_available():
+                skip_message = "Skipping the test in XPU"
+                if message:
+                    skip_message += f": {message}"
+                if has_pytest:
+                    pytest.skip(skip_message)
+                else:
+                    unittest.skip(skip_message)
+            return func(*args, **kwargs)
 
-        return unittest.skipIf(torch.xpu.is_available(), reason)(func)
+        return wrapper
 
+    # Handle both @skip_if_xpu and @skip_if_xpu() syntax
+    if callable(message):
+        func = message
+        message = None
+        return decorator(func)
     return decorator