Skip to content

GitLab

Unverified Commit 1e4a1913 authored by Thomas Wolf's avatar Thomas Wolf Committed by GitHub
Browse files

Merge pull request #1284 from slayton58/pooler_end_logits_fp16_fix 

Fix fp16 masking in PoolerEndLogits
parents 6971556a c50783e3
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 1 deletion
transformers/modeling_utils.py
View file @ 1e4a1913
...@@ -501,6 +501,9 @@ class PoolerEndLogits(nn.Module): ...@@ -501,6 +501,9 @@ class PoolerEndLogits(nn.Module):
x = self.dense_1(x).squeeze(-1) x = self.dense_1(x).squeeze(-1)
if p_mask is not None: if p_mask is not None:
if next(self.parameters()).dtype == torch.float16:
x = x * (1 - p_mask) - 65500 * p_mask
else:
x = x * (1 - p_mask) - 1e30 * p_mask x = x * (1 - p_mask) - 1e30 * p_mask
return x return x
......
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