test_fsdp_precision_forward_context_bf16

Author: Borda

tests/tests_pytorch/plugins/precision/test_fsdp.py

@@ -74,8 +74,8 @@ def test_fsdp_precision_scaler_with_bf16():
 
 
 @RunIf(min_cuda_gpus=1)
-def test_fsdp_precision_forward_context():
-    """Test to ensure that the context manager correctly is set to bfloat16."""
+def test_fsdp_precision_forward_context_f16():
+    """Test to ensure that the context manager correctly is set to float16."""
     from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
 
     precision = FSDPPrecision(precision="16-mixed")
@@ -94,6 +94,10 @@ def test_fsdp_precision_forward_context():
     assert isinstance(precision.forward_context(), _DtypeContextManager)
     assert precision.forward_context()._new_dtype == torch.float16
 
+
+@RunIf(min_cuda_gpus=1, bf16_cuda=True)
+def test_fsdp_precision_forward_context_bf16():
+    """Test to ensure that the context manager correctly is set to bfloat16."""
     precision = FSDPPrecision(precision="bf16-mixed")
     assert precision.scaler is None
     with precision.forward_context():