def test_init_module_context(monkeypatch):
    """Test that the strategy returns the context manager for initializing the module."""
    trainer = Trainer(accelerator="cpu", devices=1)
    strategy = SingleDeviceStrategy(device=torch.device("cuda"))
    strategy.tensor_init_context = Mock(wraps=strategy.tensor_init_context)
    trainer._accelerator_connector.strategy = strategy
    with trainer.init_module():
        pass
    strategy.tensor_init_context.assert_called_once_with(empty_init=None)
    strategy.tensor_init_context.reset_mock()
Test that the strategy returns the context manager for initializing the module.
tests/tests_pytorch/trainer/test_trainer.py