I’ll post lines that are correlated to this question below.
device = "cuda"

model = nn.DataParallel(model)  # For multi-gpu
model = model.module  # For multi-gpu
model.to(device)

optimizer_1 = optim.Adam([{'params': model.shared.parameters()},
                              {'params': model.unique_a.parameters()}],
                             lr=5e-5
                             )
optimizer_2 = optim.Adam([{'params': model.encoder.parameters()},
                              {'params': model.unique_b.parameters()}],
                             lr=5e-5
                             )
criterion = nn.L1Loss()
epochs = 30
for epoch in range(epochs):
        output = model(...)
        loss = criterion(output, target)
        loss.backward()
        optimizer_1.step()
        print("A Outside: output size", output.shape)

Within the model I have a print statement to print “Inside”, and I wish to see “inside” printed twice and “outside” printed once as I have 2 GPUs.
Please note: when I comment out anything related to optimizers and loss, I got the output I wanted. However, when optimizers and loss are included, it runs but I can easily see the output to show that it was only utilizing one GPU.
Please let me know if you need me to post any other piece of code.
Thanks!