gpus > 1

[hosseinfani]

Seems allRank suffers from what has been raised in this link

https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/3

So, instead of directly load the save model by:
model.load_state_dict(load_state_dict_from_file('model.pkl', dev))

We have to do the following:

state_dict = load_state_dict_from_file('model.pkl', dev)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] if 'module.' in k else k
new_state_dict[name] = v
model.load_state_dict(new_state_dict)

[mhsyno]

Hello, thank you for pointing that out. We will use this fix in the next coming PR.