clean up apex integration

2018-12-13 13:02:17 +01:00
parent 4946c2c500
commit 52c53f39d0
3 changed files with 5 additions and 5 deletions
--- a/pytorch_pretrained_bert/modeling.py
+++ b/pytorch_pretrained_bert/modeling.py
@@ -516,9 +516,9 @@ class PreTrainedBertModel(nn.Module):
        for key in state_dict.keys():
            new_key = None
            if 'gamma' in key:
-                new_key = key.replace('gamma','weight')
+                new_key = key.replace('gamma', 'weight')
            if 'beta' in key:
-                new_key = key.replace('beta','bias')
+                new_key = key.replace('beta', 'bias')
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)