Skip to content

Commit 8d3a471

Browse files
committed
add some useful decorators
1 parent c0c83fb commit 8d3a471

1 file changed

Lines changed: 23 additions & 0 deletions

File tree

src/accelerate/decorator.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
from functools import wraps
2+
3+
4+
def on_main_process(func):
5+
"""
6+
Run func on main process only
7+
"""
8+
@wraps(func)
9+
def wrapper(self, *args, **kwargs):
10+
if self.is_main_process or not self.use_distributed:
11+
return func(self, *args, **kwargs)
12+
return wrapper
13+
14+
15+
def on_local_main_process(func):
16+
"""
17+
Run func on local main process only
18+
"""
19+
@wraps(func)
20+
def wrapper(self, *args, **kwargs):
21+
if self.is_local_main_process or not self.use_distributed:
22+
return func(self, *args, **kwargs)
23+
return wrapper

0 commit comments

Comments
 (0)