susnato commited on
Commit
65eccaf
1 Parent(s): 8694ef2

Update pytorch_weights_postprocess.py

Browse files
Files changed (1) hide show
  1. pytorch_weights_postprocess.py +2 -2
pytorch_weights_postprocess.py CHANGED
@@ -55,8 +55,8 @@ for i in range(12):
55
 
56
  #Pooler
57
  for k, v in full_state_dict.items():
58
- if "pooler" in k:
59
- mod_dict.update({k:v})
60
 
61
 
62
  for k, v in mod_dict.items():
 
55
 
56
  #Pooler
57
  for k, v in full_state_dict.items():
58
+ if "dense" in k:
59
+ mod_dict.update({f"pooler.{k}":v})
60
 
61
 
62
  for k, v in mod_dict.items():