code: def test_module_init_context(
      device, precision, dtype, empty_init, monkeypatch
  ):
      """Test that the module under the init-module-context gets moved to the right device and dtype."""
      init_mock = Mock()
      monkeypatch.setattr(torch.Tensor, "uniform_", init_mock)
  
      device = torch.device(device)
      strategy = SingleDeviceStrategy(
          device=device, precision_plugin=precision
      )  # surrogate class to test base class
      with strategy.tensor_init_context(empty_init=empty_init):
          module = torch.nn.Linear(2, 2)
  
      expected_device = (
          device if _TORCH_GREATER_EQUAL_2_0 else torch.device("cpu")
      )
      assert (
          module.weight.device == module.bias.device == expected_device
      )
      assert module.weight.dtype == module.bias.dtype == dtype
      if not empty_init:
          init_mock.assert_called()
      else:
          init_mock.assert_not_called()
  
file_path: pytorch-lightning/tests/tests_pytorch/strategies/test_common.py
imports: from unittest.mock import Mock
  import pytest
  import torch
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch import Trainer
  from lightning.pytorch.plugins import DoublePrecision, HalfPrecision, Precision
  from lightning.pytorch.strategies import SingleDeviceStrategy
  from tests_pytorch.helpers.datamodules import ClassifDataModule
  from tests_pytorch.helpers.runif import RunIf
  from tests_pytorch.helpers.simple_models import ClassificationModel
label_desc: Use image node_id 2 for calling a global function with example usage: test_module_init_context(device, precision, dtype, empty_init, monkeypatch) without return types