Unverified Commit c52d262c authored by Gustaf Ahdritz's avatar Gustaf Ahdritz Committed by GitHub
Browse files

Prevent division by 0

parent b5fa2ba3
...@@ -395,8 +395,11 @@ class TriangleMultiplicativeUpdate(nn.Module): ...@@ -395,8 +395,11 @@ class TriangleMultiplicativeUpdate(nn.Module):
# Prevents overflow of torch.matmul in combine projections in # Prevents overflow of torch.matmul in combine projections in
# reduced-precision modes # reduced-precision modes
a = a / a.std() a_std = a.std()
b = b / b.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()): if(is_fp16_enabled()):
with torch.cuda.amp.autocast(enabled=False): with torch.cuda.amp.autocast(enabled=False):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment