MCPcopy Create free account
hub / github.com/NVIDIA/FasterTransformer / fix_key

Function fix_key

examples/pytorch/decoding/decoding_example.py:111–116  ·  view source on GitHub ↗
(s)

Source from the content-addressed store, hash-verified

109 ckpt = torch.load('./pytorch/translation/models/averaged-10-epoch.pt')
110 import re
111 def fix_key(s):
112 s = re.sub(r'(.*)\.layer_norm((_\d+)?)\.b_2',
113 r'\1.layer_norm\2.bias', s)
114 s = re.sub(r'(.*)\.layer_norm((_\d+)?)\.a_2',
115 r'\1.layer_norm\2.weight', s)
116 return s
117 ckpt['model'] = {fix_key(k): v for k, v in ckpt['model'].items()}
118 weights = DecodingWeights(layer_num, hidden_dim, vocab_size, ckpt)
119 else:

Callers 1

mainFunction · 0.70

Calls

no outgoing calls

Tested by

no test coverage detected