Source code for packnet_sfm.utils.horovod


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)
########################################################################################################################