added xla_device_utils to test

pytorch_lightning/utilities/xla_device_utils.py

@@ -21,13 +21,14 @@
 TORCHXLA_AVAILABLE = importlib.util.find_spec("torch_xla") is not None
 if TORCHXLA_AVAILABLE:
     import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.xla_multiprocessing as xmp
 else:
     xm = None
 
 
-def inner_f(queue, func, *args, **kwargs):  # pragma: no cover
+def inner_f(_, queue, func, *args):  # pragma: no cover
     try:
-        queue.put(func(*args, **kwargs))
+        queue.put(func(*args))
     except Exception:
         import traceback
 
@@ -39,13 +40,13 @@ def pl_multi_process(func):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         queue = Queue()
-        proc = Process(target=inner_f, args=(queue, func, *args), kwargs=kwargs)
-        proc.start()
-        proc.join(10)
-        try:
-            return queue.get_nowait()
-        except q.Empty:
-            return False
+        xmp.spawn(inner_f,
+          args=(queue, func, *args),
+          nprocs=1,
+          join=True,
+          daemon=False,
+          start_method='fork')
+        return queue.get()
 
     return wrapper