mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-08-26 11:24:41 +08:00
58a467dd68
* try delete in async or DDP us0-ecase * changelog * add model chekpoint rank * simple delete * flake8 * use global rank * chnagelog * fix review * fix import * proposal * proposal * proposal * improve proposal (fix problems with method call self) * cleaning Co-authored-by: Adrian Wälchli <adrian.waelchli@students.unibe.ch> Co-authored-by: William Falcon <waf2107@columbia.edu>
27 lines
517 B
Python
27 lines
517 B
Python
from functools import wraps
|
|
import warnings
|
|
|
|
|
|
def rank_zero_only(fn):
|
|
|
|
@wraps(fn)
|
|
def wrapped_fn(*args, **kwargs):
|
|
if rank_zero_only.rank == 0:
|
|
return fn(*args, **kwargs)
|
|
|
|
return wrapped_fn
|
|
|
|
|
|
try:
|
|
# add the attribute to the function but don't overwrite in case Trainer has already set it
|
|
getattr(rank_zero_only, 'rank')
|
|
except AttributeError:
|
|
rank_zero_only.rank = 0
|
|
|
|
|
|
def _warn(*args, **kwargs):
|
|
warnings.warn(*args, **kwargs)
|
|
|
|
|
|
rank_zero_warn = rank_zero_only(_warn)
|