mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-22 13:30:11 +08:00
Jirka Borovec
17f58d2e11
* add rank warning * changelog * use rank_zero_warn * user trainer_init * replace warnings * fix test * flake8 * docs * changelog * bug lol
19 lines
371 B
Python
19 lines
371 B
Python
"""Custom Lightning warnings"""
|
|
|
|
import warnings
|
|
|
|
_proc_rank = 0
|
|
|
|
|
|
def set_proc_rank(value: int) -> None:
|
|
"""Set the (sub)process rank."""
|
|
global _proc_rank
|
|
_proc_rank = value
|
|
|
|
|
|
def rank_zero_warn(*args, **kwargs) -> None:
|
|
"""Warning only if (sub)process has rank 0."""
|
|
global _proc_rank
|
|
if _proc_rank == 0:
|
|
warnings.warn(*args, **kwargs)
|