mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
3bf922a6ce
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/106249 Approved by: https://github.com/Skylion007
34 lines
762 B
Python
34 lines
762 B
Python
import sys
|
|
|
|
__all__ = ["register_after_fork"]
|
|
|
|
if sys.platform == "win32":
|
|
import multiprocessing.util as _util
|
|
|
|
def _register(func):
|
|
def wrapper(arg):
|
|
func()
|
|
|
|
_util.register_after_fork(_register, wrapper)
|
|
|
|
else:
|
|
import os
|
|
|
|
def _register(func):
|
|
os.register_at_fork(after_in_child=func)
|
|
|
|
|
|
def register_after_fork(func):
|
|
"""Register a callable to be executed in the child process after a fork.
|
|
|
|
Note:
|
|
In python < 3.7 this will only work with processes created using the
|
|
``multiprocessing`` module. In python >= 3.7 it also works with
|
|
``os.fork()``.
|
|
|
|
Args:
|
|
func (function): Function taking no arguments to be called in the child after fork
|
|
|
|
"""
|
|
_register(func)
|