Shared variable in GPU

[mrshenli (Shen Li)]

This should be doable, see the example code in this post:

`Exception: process 0 terminated with exit code 1` error when using `torch.multiprocessing.spawn` to parallelize over multiple GPUs

You can use share_momery_() and torch.multiprocessing.SimpleQueue to implement IPC. E.g.: import numpy as np import torch import torch.multiprocessing as mp def func(rank, x, p2c, c2p): x_power = x.to(rank) ** rank c2p.put(x_power) # citing multiprocessing doc: Unlike CPU tensors, the # sending process is required to keep the original tensor # as long as the receiving process retains a copy of # the tensor. The refcounting is implemented under the # hood but re…