Skip to content

GitLab

Unverified Commit 7751be7c authored by theblackcat102's avatar theblackcat102 Committed by GitHub
Browse files

fix reformer apex scaling issue (#4242)

parent ac7d5f67
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
src/transformers/modeling_reformer.py
View file @ 7751be7c
...@@ -562,8 +562,8 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin): ...@@ -562,8 +562,8 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
# get correct mask values depending on precision # get correct mask values depending on precision
if query_key_dots.dtype == torch.float16: if query_key_dots.dtype == torch.float16:
self_mask_value = self.self_mask_value_float16 self_mask_value = self.self_mask_value_float16.half()
mask_value = self.mask_value_float16 mask_value = self.mask_value_float16.half()
else: else:
self_mask_value = self.self_mask_value_float32 self_mask_value = self.self_mask_value_float32
mask_value = self.mask_value_float32 mask_value = self.mask_value_float32
...@@ -834,7 +834,7 @@ class LocalSelfAttention(nn.Module, EfficientAttentionMixin): ...@@ -834,7 +834,7 @@ class LocalSelfAttention(nn.Module, EfficientAttentionMixin):
if mask is not None: if mask is not None:
# get mask tensor depending on half precision or not # get mask tensor depending on half precision or not
if query_key_dots.dtype == torch.float16: if query_key_dots.dtype == torch.float16:
mask_value = self.mask_value_float16 mask_value = self.mask_value_float16.half()
else: else:
mask_value = self.mask_value_float32 mask_value = self.mask_value_float32
......
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