Skip test when atomic operations are not supported on GPU. (#7117)

tests/python/common/ops/test_ops.py

@@ -407,20 +407,20 @@ def test_segment_mm(idtype, feat_size, dtype, tol):
 def test_gather_mm_idx_b(feat_size, dtype, tol):
     if F._default_context_str == "cpu" and dtype == torch.float16:
         pytest.skip("float16 is not supported on CPU.")
-    if (
-        F._default_context_str == "gpu"
-        and dtype == torch.bfloat16
-        and not torch.cuda.is_bf16_supported()
-    ):
+
+    if F._default_context_str == "gpu":
+        if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
             pytest.skip("BF16 is not supported.")
 
         if (
-        F._default_context_str == "gpu"
-        and dtype == torch.float16
+            dtype == torch.float16
             and torch.cuda.get_device_capability() < (7, 0)
+        ) or (
+            dtype == torch.bfloat16
+            and torch.cuda.get_device_capability() < (8, 0)
         ):
             pytest.skip(
-            f"FP16 is not supported for atomic operations on GPU with "
+                f"{dtype} is not supported for atomic operations on GPU with "
                 f"cuda capability ({torch.cuda.get_device_capability()})."
             )