ruff

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Author: akoumpa

nemo_automodel/components/distributed/parallel_styles.py

@@ -164,9 +164,7 @@ def _partition_linear_fn(self, name, module, device_mesh):
         def lora_a_output_hook(module, input, output):
             if isinstance(output, DTensor):
                 if any(isinstance(p, Shard) for p in output.placements):
-                    output = output.redistribute(
-                        device_mesh=output.device_mesh, placements=[Replicate()]
-                    )
+                    output = output.redistribute(device_mesh=output.device_mesh, placements=[Replicate()])
             return output
 
         if hasattr(module, "lora_A"):