fix grad norm fsdp

PrimeIntellect-ai · Jul 19, 2024 · e4e4f86 · e4e4f86
1 parent 1531aa5
commit e4e4f86
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/open_diloco/train_fsdp.py b/open_diloco/train_fsdp.py
@@ -402,7 +402,7 @@ def scheduler_fn(opt):
             else:
                 scaler.unscale_(optimizer=optimizer)
 
-            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # gradient clipping
+            model.clip_grad_norm_(1.0)  # gradient clipping
 
             if world_messenger_hv:
                 optimizer.step(scaler=scaler)