Clamp output when rounding weight to prevent Nan.

zcfire2017 · Oct 19, 2024 · 73e3a9e · 73e3a9e
1 parent 518c0dc
commit 73e3a9e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/comfy/float.py b/comfy/float.py
@@ -41,6 +41,8 @@ def manual_stochastic_round_to_float8(x, dtype, generator=None):
         (2.0 ** (-EXPONENT_BIAS + 1)) * abs_x
     )
 
+    inf = torch.finfo(dtype)
+    torch.clamp(sign, min=inf.min, max=inf.max, out=sign)
     return sign