SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

fix reformer apex scaling issue (#4242)

Browse Source
This commit is contained in:
theblackcat102
2020-05-11 22:53:42 +08:00
committed by GitHub
parent ac7d5f67a2
commit 7751be7cee
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/transformers/modeling_reformer.py
View File

@@ -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):
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