Hi,
What are you trying to achieve here, train 2 models (a and b) at the same time? You generate value f every time for node 1 and node 2 (f can differ between 1 and 2, right? so the code path chosen by node 1 and 2 would be different) and then you are trying to sync the gradients (via DDP). If so, node 1 can run
out_b = ddp_netb(torch.randn(0,3, 32, 32).to(rank))
out_a = ddp_neta(torch.randn(6,3,128,128).to(rank))
while node 2 could run
out_a = ddp_neta(torch.randn(0,3,128,128).to(rank))
out_b = ddp_netb(torch.randn(3,3,32,32).to(rank))
for dimension with data.shape[0] == 0 you return 0 and none of parameters are being used on one node while they are being used and sync on another. This can lead to issues with DDP I believe.
Did you intend to run the same steps for each iteration of the training loop on node 0 and node 1, e.g. sync value of f between them?