Retrieval Example: Failed load from saved model

[nfrumkin]

Hello,

I was going through the retrieval example and torch.load() throws an error in two_tower_retrieval.py:

RuntimeError: Need to initialize default process group using "init_process_group" before loading ShardedTensor

I initialize the process group as the error indicates and run into the following error on dmp.load_state_dict(retrieval_sd):

AttributeError: 'Tensor' object has no attribute 'local_shards'

To Reproduce

a. execute two_tower_train.py using command 2 in the retrieval example's README to get the retrieval model.
b. execute two_tower_retrieval.py using command 3.

Any help is appreciated!