code: def test_module_init_context(precision, expected_dtype):
      plugin = DeepSpeedPrecision(precision=precision)
      with plugin.module_init_context():
          model = torch.nn.Linear(2, 2)
          assert torch.get_default_dtype() == expected_dtype
      assert model.weight.dtype == expected_dtype
  
file_path: pytorch-lightning/tests/tests_fabric/plugins/precision/test_deepspeed.py
imports: from unittest import mock
  from unittest.mock import Mock
  import pytest
  import torch
  from lightning.fabric.plugins.precision.deepspeed import DeepSpeedPrecision
  from lightning.fabric.utilities.types import Steppable
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 6 for calling a global function with example usage: test_module_init_context(precision, expected_dtype) without return types