skip fsdp2 test

pytorch · Jan 23, 2025 · 26fe1f1 · 26fe1f1
1 parent 8fd16a6
commit 26fe1f1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
@@ -386,8 +386,10 @@ def world_size(self) -> int:
         not TORCH_VERSION_AT_LEAST_2_5, reason="PyTorch>=2.5 is required."
     )
     @skip_if_lt_x_gpu(_FSDP_WORLD_SIZE)
-    @skip_if_rocm("ROCm enablement in progress")
     def test_fsdp2(self):
+        if torch.version.hip is not None:
+            pytest.skip("ROCm enablement in progress")
+
         optim_classes = [low_bit_optim.AdamW8bit, low_bit_optim.AdamW4bit]
         if torch.cuda.get_device_capability() >= (8, 9):
             optim_classes.append(low_bit_optim.AdamWFp8)