Prevent division by 0

openfold/model/triangular_multiplicative_update.py

@@ -395,6 +395,9 @@ class TriangleMultiplicativeUpdate(nn.Module):
 
         # Prevents overflow of torch.matmul in combine projections in
         # reduced-precision modes
+        a_std = a.std()
+        b_std = b.std()
+        if(a_std != 0. and b_std != 0):
             a = a / a.std()
             b = b / b.std()