As I understand, originally the example uses 3 processes: 1 master, and worker1 + worker2, where both workers are assigned to two CUDA devices. Now you only use 2 processes: worker0 (also as master) and worker1, both are assigned to CPU, so basically worker0 runs RPC on itself.
Here are my suggestions:

I don’t know if a process runs RPC on itself will lead to memory leak or not. Ideally it should work as doing the work locally, but I am not sure if such an unreasonable corner case is supported well. You can update the code to train ResNetShard1 locally.
I think you should still use mp.spawn in the main function. I haven’t tried multi-node training by myself. Does mp.spawn fail on multiple nodes?
Out of curiosity, since I don’t have experience on running the computation on CPU cores via RPC, I wonder if you can run the same code on 2 CPU cores of the same machine, by changing “cuda:0” and “cuda:1” to “cpu:0” and “cpu:1”, respectively. We can check if this will also cause memory leak as well. If so, the problem could be TensorPipe + CPU.