Skip to content

Commit

Permalink
moved inner_f to global scope
Browse files Browse the repository at this point in the history
  • Loading branch information
lezwon committed Aug 31, 2020
1 parent ecae7d2 commit f315261
Showing 1 changed file with 11 additions and 15 deletions.
26 changes: 11 additions & 15 deletions pytorch_lightning/utilities/xla_device_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import functools
import importlib
from multiprocessing import Process, Queue

TORCHXLA_AVAILABLE = importlib.util.find_spec("torch_xla") is not None
if TORCHXLA_AVAILABLE:
Expand All @@ -8,27 +9,22 @@
xm = None


def inner_f(queue, func, **kwargs):
try:
queue.put(func(**kwargs))
except Exception:
import traceback
traceback.print_exc()
queue.put(None)


def pl_multi_process(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):

from multiprocessing import Process, Queue

queue = Queue()

def inner_f():
try:
queue.put(func(**kwargs))
except Exception:
import traceback

traceback.print_exc()
queue.put(None)

proc = Process(target=inner_f, kwargs=kwargs)
proc = Process(target=inner_f, args=(queue, func,), kwargs=kwargs)
proc.start()
proc.join()

return queue.get()

return wrapper
Expand Down

0 comments on commit f315261

Please sign in to comment.