mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
f3fce597e9
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter. You can review these PRs via: ```bash git diff --ignore-all-space --ignore-blank-lines HEAD~1 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/129769 Approved by: https://github.com/ezyang
36 lines
790 B
Python
36 lines
790 B
Python
# mypy: allow-untyped-defs
|
|
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)
|