Enable v2 forward test for ROCm (#3573)

Summary:
X-link: facebookresearch/FBGEMM#663

Since #3266 is merged, the v2 forward kernel should be tested for ROCm devices as well.

cc: liligwu

Pull Request resolved: #3573

Reviewed By: leitian, spcyppt

Differential Revision: D68237138

Pulled By: q10

fbshipit-source-id: 3f743303c13ae79976dda273ffb75deee57fd11f

fbgemm_gpu/test/tbe/training/forward_test.py

@@ -541,7 +541,7 @@ def test_forward_gpu_no_cache_fp16(
 
     @unittest.skipIf(*gpu_unavailable)
     @given(
-        use_experimental_tbe=st.booleans() if not TEST_WITH_ROCM else st.just(False),
+        use_experimental_tbe=st.booleans(),
     )
     @settings(
         verbosity=VERBOSITY,
@@ -670,7 +670,7 @@ def test_forward_gpu_uvm_cache_int8(
     @unittest.skipIf(*gpu_unavailable)
     @given(
         cache_algorithm=st.sampled_from(CacheAlgorithm),
-        use_experimental_tbe=st.booleans() if not TEST_WITH_ROCM else st.just(False),
+        use_experimental_tbe=st.booleans(),
     )
     @settings(
         verbosity=VERBOSITY,
@@ -740,7 +740,7 @@ def test_forward_gpu_uvm_cache_fp16(
     @unittest.skipIf(*gpu_unavailable)
     @given(
         cache_algorithm=st.sampled_from(CacheAlgorithm),
-        use_experimental_tbe=st.booleans() if not TEST_WITH_ROCM else st.just(False),
+        use_experimental_tbe=st.booleans(),
     )
     @settings(
         verbosity=VERBOSITY,