Merge pull request #1284 from slayton58/pooler_end_logits_fp16_fix

Fix fp16 masking in PoolerEndLogits

Merge pull request #1284 from slayton58/pooler_end_logits_fp16_fix
Fix fp16 masking in PoolerEndLogits
1e4a1913 · Thomas Wolf · GitHub · 6971556a · c50783e3 · 1e4a1913
Unverified Commit 1e4a1913 authored Oct 01, 2019 by Thomas Wolf Committed by GitHub Oct 01, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

transformers/modeling_utils.py transformers/modeling_utils.py +4 -1

No files found.
--- a/transformers/modeling_utils.py
+++ b/transformers/modeling_utils.py
@@ -501,7 +501,10 @@ class PoolerEndLogits(nn.Module):
        x = self.dense_1(x).squeeze(-1)

        if p_mask is not None:
-            x = x * (1 - p_mask) - 1e30 * p_mask
+            if next(self.parameters()).dtype == torch.float16:
+                x = x * (1 - p_mask) - 65500 * p_mask
+            else:
+                x = x * (1 - p_mask) - 1e30 * p_mask

        return x