def test_tensor_init_context(device, precision, dtype):
    """Test that tensors under the init-tensor-context get moved to the right device and dtype."""
    device = torch.device(device)
    strategy = SingleDeviceStrategy(device=device, precision=precision)  # surrogate class to test base class
    with strategy.tensor_init_context():
        tensor0 = torch.tensor(42.0)
        tensor1 = torch.tensor(42)
        tensor2 = torch.tensor(42.0, dtype=torch.half)

    assert tensor0.device == tensor1.device == tensor2.device == device
    assert tensor0.dtype == dtype
    assert tensor1.dtype == torch.long  # `.init_tensor()` only affects floating point dtypes
    assert tensor2.dtype == torch.half  # this tensor was created with an explicit dtype assignment
Test that tensors under the init-tensor-context get moved to the right device and dtype.
tests/tests_fabric/strategies/test_strategy.py