I see, I can confirm repro’ing this error case. It seems like autograd cannot handle computing gradients on tensors that have been created by operations that involve cross-process communication, so it simply refuses to serialize them for IPC in the fwd pass itself.
@pritamdamania87 Is our assessment here correct? Is this use case supported with Distributed Autograd/some other alternative? I’m guessing one way of doing this would be to send RPC’s in the fwd pass to other processes to perform the func_to_apply on some chunk of data, and then collect the results (and distributed autograd would do the reverse in the bwd pass), but not sure if this is feasible/the best approach.