Policy test

Author: carmocca

tests/tests_fabric/strategies/test_fsdp.py

@@ -26,6 +26,7 @@
 from lightning_utilities.core.imports import RequirementCache
 from torch.optim import Adam
 
+import lightning
 from lightning.fabric import Fabric
 from lightning.fabric.plugins.environments import LightningEnvironment
 from lightning.fabric.strategies import FSDPStrategy
@@ -125,12 +126,16 @@ def test_fsdp_no_backward_sync():
 
 
 @RunIf(min_torch="1.12")
-@mock.patch("lightning.fabric.strategies.fsdp._TORCH_GREATER_EQUAL_1_13", False)
-def test_fsdp_activation_checkpointing_support():
+def test_fsdp_activation_checkpointing_support(monkeypatch):
     """Test that we error out if activation checkpointing requires a newer PyTorch version."""
+    monkeypatch.setattr(lightning.fabric.strategies.fsdp, "_TORCH_GREATER_EQUAL_1_13", False)
     with pytest.raises(ValueError, match="activation_checkpointing` requires torch >= 1.13.0"):
         FSDPStrategy(activation_checkpointing=Mock())
 
+    monkeypatch.setattr(lightning.fabric.strategies.fsdp, "_TORCH_GREATER_EQUAL_2_1", False)
+    with pytest.raises(ValueError, match="activation_checkpointing_policy` requires torch >= 2.1.0"):
+        FSDPStrategy(activation_checkpointing_policy=Mock())
+
 
 @RunIf(min_torch="1.13")
 def test_fsdp_activation_checkpointing():