code: def test_fsdp_activation_checkpointing():
      """Test that the FSDP strategy can apply activation checkpointing to the given layers."""
  
      class Block1(nn.Linear):
          pass
  
      class Block2(nn.Linear):
          pass
  
      class Model(nn.Module):
          def __init__(self):
              super().__init__()
              self.layer0 = nn.Sequential(Block1(4, 4), Block1(5, 5))
              self.layer1 = Block2(2, 2)
              self.layer2 = nn.Linear(3, 3)
  
      if _TORCH_GREATER_EQUAL_2_1:
          from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  
          strategy = FSDPStrategy(
              activation_checkpointing_policy={Block1}
          )
          assert set(strategy._activation_checkpointing_kwargs) == {
              "auto_wrap_policy"
          }
          assert isinstance(
              strategy._activation_checkpointing_kwargs[
                  "auto_wrap_policy"
              ],
              ModuleWrapPolicy,
          )
  
          strategy = FSDPStrategy(
              activation_checkpointing_policy=ModuleWrapPolicy(
                  {Block1, Block2}
              )
          )
          assert set(strategy._activation_checkpointing_kwargs) == {
              "auto_wrap_policy"
          }
          assert isinstance(
              strategy._activation_checkpointing_kwargs[
                  "auto_wrap_policy"
              ],
              ModuleWrapPolicy,
          )
      else:
          strategy = FSDPStrategy(activation_checkpointing=Block1)
          assert set(strategy._activation_checkpointing_kwargs) == {
              "check_fn"
          }
  
          strategy = FSDPStrategy(
              activation_checkpointing=[Block1, Block2]
          )
          assert set(strategy._activation_checkpointing_kwargs) == {
              "check_fn"
          }
  
          strategy = FSDPStrategy(
              activation_checkpointing_policy={Block1}
          )
          assert set(strategy._activation_checkpointing_kwargs) == {
              "check_fn"
          }
  
          strategy = FSDPStrategy(
              activation_checkpointing_policy={Block1, Block2}
          )
          assert set(strategy._activation_checkpointing_kwargs) == {
              "check_fn"
          }
  
      strategy._parallel_devices = [torch.device("cuda", 0)]
      with mock.patch(
          "torch.distributed.fsdp.FullyShardedDataParallel",
          new=MagicMock,
      ), mock.patch(
          "torch.distributed.algorithms._checkpoint.checkpoint_wrapper.apply_activation_checkpointing"
      ) as apply_mock:
          wrapped = strategy.setup_module(Model())
      apply_mock.assert_called_with(
          wrapped,
          checkpoint_wrapper_fn=ANY,
          **strategy._activation_checkpointing_kwargs,
      )
  
imports: from datetime import timedelta
  from re import escape
  from unittest import mock
  from unittest.mock import ANY, MagicMock, Mock
  import lightning.fabric
  import pytest
  import torch
  import torch.nn
  from lightning.fabric.plugins import HalfPrecision
  from lightning.fabric.plugins.environments import LightningEnvironment
  from lightning.fabric.strategies import FSDPStrategy
  from lightning.fabric.strategies.fsdp import _FSDPBackwardSyncControl, _get_full_state_dict_context, _has_meta_device_parameters, _is_sharded_checkpoint
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
  from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, FullyShardedDataParallel, MixedPrecision
  from torch.optim import Adam
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 10 for calling a global function with example usage: test_fsdp_activation_checkpointing() without return types