code: def test_module_init_context(precision, expected_dtype):
      """Test that the module under the init-context gets moved to the right device and dtype."""
      fabric = Fabric(
          accelerator="cuda",
          devices=2,
          strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
          precision=precision,
      )
      fabric.launch()
  
      def _run_setup_assertions(empty_init, expected_device):
          with fabric.init_module(empty_init=empty_init):
              model = torch.nn.Linear(100, 100, bias=False)
  
          # The model is on the CPU/meta-device until after `.setup()``
          assert model.weight.device == expected_device
          assert model.weight.dtype == expected_dtype
          model = fabric.setup(model)
          # Parameters get sharded in `.setup()` and moved to the target device
          assert model.weight.device == torch.device(
              "cuda", fabric.local_rank
          )
          assert model.weight.dtype == expected_dtype
  
      # Case 1: No empty init
      _run_setup_assertions(
          empty_init=False, expected_device=torch.device("cpu")
      )
  
      if _TORCH_GREATER_EQUAL_2_1:
          # Case 2: Empty-init with PyTorch >= 2.1 supports meta device
          _run_setup_assertions(
              empty_init=True, expected_device=torch.device("meta")
          )
      else:
          # Case 2: Empty-init with PyTorch < 2.1 only supports `torch.empty()`-init
          _run_setup_assertions(
              empty_init=True, expected_device=torch.device("cpu")
          )
  
imports: import os
  from copy import deepcopy
  from pathlib import Path
  from unittest import mock
  import pytest
  import torch
  from lightning.fabric import Fabric
  from lightning.fabric.plugins import FSDPPrecision
  from lightning.fabric.strategies import FSDPStrategy
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0, _TORCH_GREATER_EQUAL_2_1
  from lightning.fabric.wrappers import _FabricOptimizer
  from torch.distributed.fsdp import FlatParameter, FullyShardedDataParallel, OptimStateKeyType
  from torch.distributed.fsdp.wrap import always_wrap_policy, wrap
  from torch.nn import Parameter
  from tests_fabric.helpers.models import BoringFabric
  from tests_fabric.helpers.runif import RunIf
  from tests_fabric.test_fabric import BoringModel
label_desc: Use image node_id 7 for calling a global function with example usage: test_module_init_context(precision, expected_dtype) without return types