From ee282be91fb0da72e6057ed5ad371631ca601fa2 Mon Sep 17 00:00:00 2001
From: rockerBOO <rockerboo@gmail.com>
Date: Thu, 19 Jun 2025 16:28:18 -0400
Subject: [PATCH] Revert zero_grad

---
 train_network.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/train_network.py b/train_network.py
index 77cbe211..af13135d 100644
--- a/train_network.py
+++ b/train_network.py
@@ -1435,8 +1435,7 @@ class NetworkTrainer:
 
                     optimizer.step()
                     lr_scheduler.step()
-                    # optimizer.zero_grad(set_to_none=True)
-                    optimizer.zero_grad(set_to_none=False)
+                    optimizer.zero_grad(set_to_none=True)
 
                 if args.scale_weight_norms:
                     keys_scaled, mean_norm, maximum_norm = accelerator.unwrap_model(network).apply_max_norm_regularization(