def test_amp_precision_forward_context():
    """Test to ensure that the context manager correctly is set to bfloat16 on CPU and CUDA."""
    precision = MixedPrecision(precision="16-mixed", device="cuda")
    assert precision.device == "cuda"
    scaler_cls = torch.amp.GradScaler if _TORCH_GREATER_EQUAL_2_4 else torch.cuda.amp.GradScaler
    assert isinstance(precision.scaler, scaler_cls)
    assert torch.get_default_dtype() == torch.float32
    with precision.forward_context():
        assert torch.get_autocast_gpu_dtype() == torch.float16

    precision = MixedPrecision(precision="bf16-mixed", device="cpu")
    assert precision.device == "cpu"
    assert precision.scaler is None
    with precision.forward_context():
        assert torch.get_autocast_cpu_dtype() == torch.bfloat16

    context_manager = precision.forward_context()
    assert isinstance(context_manager, torch.autocast)
    assert context_manager.fast_dtype == torch.bfloat16
Test to ensure that the context manager correctly is set to bfloat16 on CPU and CUDA.
tests/tests_fabric/plugins/precision/test_amp.py