From 3ee548cd06445d559e23e72942dea9675df5058a Mon Sep 17 00:00:00 2001
From: Benoit <69694610+BenoitDalFerro@users.noreply.github.com>
Date: Thu, 19 Sep 2024 19:18:55 +0200
Subject: [PATCH] Update AdEMAMix-Shampoo.py

Fixing division by 0 error at first step
---
 AdEMAMix-Shampoo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/AdEMAMix-Shampoo.py b/AdEMAMix-Shampoo.py
index 2a73931..d2e5dfb 100644
--- a/AdEMAMix-Shampoo.py
+++ b/AdEMAMix-Shampoo.py
@@ -284,7 +284,7 @@ def _update_adamemix_distributed_shampoo(
             denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
 
             # Compute step size
-            step_size = lr / bias_correction1
+            step_size = lr / (bias_correction1 if bias_correction1 > 0 else 0.01)
 
             # Apply weight decay
             if weight_decay != 0: