Refine _platform_supports_mx_gemm check

Signed-off-by: Jagadish Krishnamoorthy <[email protected]>

Author: jagadish-amd

torch/testing/_internal/common_cuda.py

@@ -105,10 +105,11 @@ def evaluate_platform_supports_fp8():
 PLATFORM_SUPPORTS_FP8: bool = LazyVal(lambda: evaluate_platform_supports_fp8())
 
 def _platform_supports_mx_gemm():
-    if TEST_CUDA:
-        return SM100OrLater
-    if TEST_WITH_ROCM:
-        return torch.cuda.get_device_properties(torch.cuda.current_device(0)).name.startswith('gfx950')
+    if torch.cuda.is_available():
+        if torch.version.hip:
+            return 'gfx95' in torch.cuda.get_device_properties(0).gcnArchName
+        else:
+            return SM100OrLater
     return False
 
 PLATFORM_SUPPORTS_MX_GEMM: bool = LazyVal(lambda: _platform_supports_mx_gemm())