mypy

Author: carmocca

src/lightning/fabric/strategies/fsdp.py

@@ -621,10 +621,10 @@ def _activation_checkpointing_kwargs(
                 f"({set(classes)}))` instead."
             )
         return {"check_fn": lambda submodule: isinstance(submodule, classes)}
-    if activation_checkpointing_policy is not None:
-        if not _TORCH_GREATER_EQUAL_2_1:
-            raise ValueError("`activation_checkpointing_policy` requires torch >= 2.1.0. HINT: `pip install -U torch`")
-        return {"auto_wrap_policy": activation_checkpointing_policy}
+    assert activation_checkpointing_policy is not None
+    if not _TORCH_GREATER_EQUAL_2_1:
+        raise ValueError("`activation_checkpointing_policy` requires torch >= 2.1.0. HINT: `pip install -U torch`")
+    return {"auto_wrap_policy": activation_checkpointing_policy}
 
 
 def _setup_activation_checkpointing(module: Module, activation_checkpointing_kwargs: Dict) -> None: