Prevent division by 0

c52d262c · Gustaf Ahdritz · GitHub · b5fa2ba3 · c52d262c
Unverified Commit c52d262c authored May 13, 2023 by Gustaf Ahdritz Committed by GitHub May 13, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 2 deletions

openfold/model/triangular_multiplicative_update.py openfold/model/triangular_multiplicative_update.py +5 -2

No files found.
--- a/openfold/model/triangular_multiplicative_update.py
+++ b/openfold/model/triangular_multiplicative_update.py
@@ -395,8 +395,11 @@ class TriangleMultiplicativeUpdate(nn.Module):

        # Prevents overflow of torch.matmul in combine projections in
        # reduced-precision modes
-        a = a / a.std()
-        b = b / b.std()
+        a_std = a.std()
+        b_std = b.std()
+        if(a_std != 0. and b_std != 0):
+            a = a / a.std()
+            b = b / b.std()

        if(is_fp16_enabled()):
            with torch.cuda.amp.autocast(enabled=False):