import horovod.torch as hvd
########################################################################################################################
[docs]def hvd_init():
hvd.init()
[docs]def on_rank_0(func):
def wrapper(*args, **kwargs):
if rank() == 0:
func(*args, **kwargs)
return wrapper
[docs]def rank():
return hvd.rank()
[docs]def world_size():
return hvd.size()
[docs]@on_rank_0
def print0(string='\n'):
print(string)
########################################################################################################################