disable use_cpu test (#1635)

Summary:
This PR addresses the issue #1636

akin to https://github.com/pytorch/FBGEMM/blob/8616ed701015f8b9e4c2825ce592b204b4cfaf28/fbgemm_gpu/test/split_table_batched_embeddings_test.py#L1009

Pull Request resolved: #1635

Reviewed By: shintaro-iwasaki

Differential Revision: D44033725

Pulled By: q10

fbshipit-source-id: 49f28fc2f1c20948a42728eebf3defc5195baa5d

fbgemm_gpu/test/jagged_tensor_ops_test.py

@@ -20,14 +20,20 @@
     from fbgemm_gpu import open_source  # noqa: F401
 
     # pyre-ignore[21]
-    from test_utils import gpu_available, gpu_unavailable, running_on_github
+    from test_utils import (
+        gpu_available,
+        gpu_unavailable,
+        running_on_github,
+        TEST_WITH_ROCM,
+    )
 except Exception:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
     from fbgemm_gpu.test.test_utils import (
         gpu_available,
         gpu_unavailable,
         running_on_github,
+        TEST_WITH_ROCM,
     )
 
 
@@ -1466,7 +1472,11 @@ def jagged_index_select_2d_ref(
                 torch.long,
             ]  # Disable torch.bfloat16 due to large error bound
         ),
-        use_cpu=st.booleans() if gpu_available else st.just(True),
+        use_cpu=st.booleans()
+        if (gpu_available and not TEST_WITH_ROCM)
+        else st.just(False)
+        if (gpu_available and TEST_WITH_ROCM)
+        else st.just(True),
     )
     @settings(max_examples=20, deadline=None)
     def test_jagged_index_select_2d(