[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

tests/pytorch/distributed/run_fsdp2_fused_adam.py

@@ -122,14 +122,9 @@ def test_fused_adam_fp8_master_weights(recipe=None):
     model = _build_model(fp8_init=True, recipe=recipe)
 
     # Verify FP8 params created
-    qt_count = sum(
-        1
-        for _, p in model.named_parameters()
-        if isinstance(p, QuantizedTensor)
-    )
+    qt_count = sum(1 for _, p in model.named_parameters() if isinstance(p, QuantizedTensor))
     assert qt_count > 0, "No QuantizedTensor local tensors before training"
 
-
     model = _shard_model(model, world_size)
 
     # Verify params are DTensors
@@ -144,7 +139,6 @@ def test_fused_adam_fp8_master_weights(recipe=None):
     )
     assert qt_count > 0, "No QuantizedTensor local tensors after sharding"
 
-
     optimizer = te.optimizers.FusedAdam(
         model.parameters(),
         lr=1e-3,

tests/pytorch/distributed/run_fsdp2_model.py

@@ -328,7 +328,11 @@ def _train(args):
         target = torch.randn(out_shape, device=device)
 
         # NVFP4BlockScaling requires bfloat16 inputs in both the forward and backward passes.
-        with torch.autocast(device_type="cuda", dtype=torch.bfloat16) if args.recipe == "NVFP4BlockScaling" else nullcontext():
+        with (
+            torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+            if args.recipe == "NVFP4BlockScaling"
+            else nullcontext()
+        ):
             with te.autocast(enabled=True, recipe=fp8_recipe):
                 output = model(input_data)
                 loss = F.mse_loss(output, target)

tests/pytorch/distributed/test_torch_fsdp2.py

@@ -73,9 +73,7 @@ def _run_test(fp_init, sharding_dims, recipe, layer_type):
 def test_distributed(fp8_init, sharding_dims, fp_recipe, layer_type):
 
     if fp_recipe in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init:
-        pytest.xfail(
-            f"{fp_recipe} + fp8_init: test_fp8_fsdp2_allgather is currently failing."
-        )
+        pytest.xfail(f"{fp_recipe} + fp8_init: test_fp8_fsdp2_allgather is currently failing.")
 
     _run_test(fp8_init, sharding_dims, fp_recipe, layer_type)