Unexpected difference torch.multiprocessing.manager.queue and torch.multiprocessing.queue

[shayben]

It seems that torch.multiprocessing.Manager.Queue doesn't support multiprocessor sharing of cuda tensors.
Running the repro below throws the exception:

RuntimeError: Attempted to send CUDA tensor received from another process; this is not currently supported. Consider cloning before sending.

import torch
import torch.multiprocessing as mp


def produce(in_Q, out_Q):
    while True:
        t = in_Q.get()
        v = torch.tensor(t, device='cuda')
        out_Q.put(v)

def consume(out_Q):
    res = out_Q.get()
    return res

def main():
    if False: #-----------Change this to True for the expected behavior queue --------------------
        qsrc = mp
    else:
        manager = mp.Manager()
        qsrc = manager
    
    in_Q = qsrc.Queue(10)
    out_Q = qsrc.Queue(10)

    consumer = mp.Process(target=consume, args=[out_Q])
    consumer.start()
    producer = mp.Process(target=produce, args=(in_Q, out_Q))
    producer.start()
    in_Q.put(1)
    consumer.join()
    print('done.')

if __name__ == "__main__":
    main()

cc @ngimel

[michelgokan]

Any suggestions on possible alternatives?

[berooo]

have you solved it?

[marvelous-melanie]

Hi, I'm wondering if there are any updates on this issue?

[AOTY-szy]

have you solved it?

[btalberg]

I'm struggling with the same issue using torch 2.1.0