Hi! Is this the only adjustment to make DDP work with two separate models?
I am trying to train two models in two stages: I train the first DDP model A for k steps, then I use its encoder part with torch.no_grad() and then use its projection layers and a second DDP model B.
For this I use two data loaders with distributed data samplers and two optimizers, each watching only one model.
So in the training loop I have:
for k steps:
 pred = modelA(X)
 loss = lossA(pred, y)
 loss.backward()
 optimizerA.step()
....
for m steps:
 with torch.no_grad():
      x = modelA.module.encoder()
 x = modelA.module.proj(x)
 pred = modelB(x)
 loss = lossB(pred, y)
 loss.backward()
 optimizerA.step()
 optimizerB.step()

If I do not use 2 pg I get a CUDA memory access problems as you said, but with 2 pgs it just seems to hang after training modelA for some amount of steps
What should I do?