copied pl_multi_process

pytorch_lightning/utilities/xla_device_utils.py

@@ -19,13 +19,29 @@ def inner_f(queue, func, **kwargs):
 
 
 def pl_multi_process(func):
+    """Wrapper for running multi-processing tests."""
+
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
+
+        from multiprocessing import Process, Queue
         queue = Queue()
-        proc = Process(target=inner_f, args=(queue, func,), kwargs=kwargs)
+
+        def inner_f(queue, **kwargs):
+            try:
+                result = func(**kwargs)
+                queue.put(result)
+            except Exception:
+                import traceback
+                traceback.print_exc()
+                queue.put(False)
+
+        proc = Process(target=inner_f, args=(queue,), kwargs=kwargs)
         proc.start()
         proc.join()
-        return queue.get()
+
+        result = queue.get()
+        assert result == 1, 'expected 1, but returned %s' % result
 
     return wrapper