def test_cpu_amp_precision_context_manager():
    """Test to ensure that the context manager correctly is set to CPU + bfloat16."""
    plugin = MixedPrecision("bf16-mixed", "cpu")
    assert plugin.device == "cpu"
    assert plugin.scaler is None
    context_manager = plugin.autocast_context_manager()
    assert isinstance(context_manager, torch.autocast)
    assert context_manager.fast_dtype == torch.bfloat16
Test to ensure that the context manager correctly is set to CPU + bfloat16.
tests/tests_pytorch/plugins/test_amp_plugins.py