Skip to content

Commit

Permalink
Fix race condiction (#12)
Browse files Browse the repository at this point in the history
Fix #4
  • Loading branch information
yjc4 committed Nov 29, 2022
1 parent bfd0ffd commit 901f860
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion basic_utils.py
Expand Up @@ -72,6 +72,7 @@ def load_model_emb(args, tokenizer):
### random emb or pre-defined embedding like glove embedding. You can custome your own init here.
model = torch.nn.Embedding(tokenizer.vocab_size, args.hidden_dim)
path_save = '{}/random_emb.torch'.format(args.checkpoint_path)
path_save_ind = path_save + ".done"
if int(os.environ['LOCAL_RANK']) == 0:
if os.path.exists(path_save):
print('reload the random embeddings', model)
Expand All @@ -80,8 +81,11 @@ def load_model_emb(args, tokenizer):
print('initializing the random embeddings', model)
torch.nn.init.normal_(model.weight)
torch.save(model.state_dict(), path_save)
os.sync()
with open(path_save_ind, "x") as _:
pass
else:
while not os.path.exists(path_save):
while not os.path.exists(path_save_ind):
time.sleep(1)
print('reload the random embeddings', model)
model.load_state_dict(torch.load(path_save))
Expand Down

0 comments on commit 901f860

Please sign in to comment.