Skip to content

Commit 398f4e8

Browse files
authored
Wrap horovod code inside main function (#739)
1 parent a3f4f4b commit 398f4e8

File tree

1 file changed

+9
-5
lines changed

1 file changed

+9
-5
lines changed

tests/data/horovod/test_hvd_basic.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
import json
22
import os
3+
34
import horovod.tensorflow as hvd
45

5-
hvd.init()
6+
if __name__ == '__main__':
7+
8+
hvd.init()
9+
10+
with open(os.path.join('/opt/ml/model/rank-%s' % hvd.rank()), 'w+') as f:
11+
basic_info = {'rank': hvd.rank(), 'size': hvd.size()}
612

7-
with open(os.path.join('/opt/ml/model/rank-%s' % hvd.rank()), 'w+') as f:
8-
basic_info = {'rank': hvd.rank(), 'size': hvd.size()}
13+
json.dump(basic_info, f)
14+
print('Saved file "rank-%s": %s' % (hvd.rank(), basic_info))
915

10-
print(basic_info)
11-
json.dump(basic_info, f)

0 commit comments

Comments
 (0)