code: def test_tensor_init_context(precision_plugin, expected_dtype):
      """Test that the module under the init-context gets moved to the right device and dtype."""
      parallel_devices = [
          torch.device("cuda", 0),
          torch.device("cuda", 1),
      ]
      expected_device = (
          parallel_devices[1]
          if _TORCH_GREATER_EQUAL_2_0
          else torch.device("cpu")
      )
  
      strategy = DDPStrategy(
          parallel_devices=parallel_devices,
          precision_plugin=precision_plugin,
          cluster_environment=LightningEnvironment(),
      )
      assert strategy.local_rank == 1
      with strategy.tensor_init_context():
          module = torch.nn.Linear(2, 2)
      assert (
          module.weight.device == module.bias.device == expected_device
      )
      assert module.weight.dtype == module.bias.dtype == expected_dtype
  
file_path: pytorch-lightning/tests/tests_pytorch/strategies/test_ddp.py
imports: import os
  from datetime import timedelta
  from unittest import mock
  import pytest
  import torch
  from lightning.fabric.plugins.environments import LightningEnvironment
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch import LightningModule, Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from lightning.pytorch.plugins import DoublePrecision, HalfPrecision, Precision
  from lightning.pytorch.strategies import DDPStrategy
  from lightning.pytorch.trainer.states import TrainerFn
  from torch.nn.parallel import DistributedDataParallel
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 3 for calling a global function with example usage: test_tensor_init_context(precision_plugin, expected_dtype) without return types