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)

    assert module.weight.device == module.bias.device == device
    assert module.weight.dtype == module.bias.dtype == dtype
    if not empty_init:
        init_mock.assert_called()
    else:
        init_mock.assert_not_called()
Test that the module under the init-module-context gets moved to the right device and dtype.
tests/tests_pytorch/strategies/test_common.py