Hi albanD,
Thank you for your reply. I do use dataparallel with:
model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()
And this is the code snippet for computing loss and doing backward prop:
out_dict = model(torch.cat((s_input_var, t_input_var), dim=0))
output, ad_out = out_dict['output'], out_dict['ad_out']
split_size = args.batch_size // len(args.gpus)
ad_label = []
for i in range(len(args.gpus)):
    ad_label.append(torch.zeros([split_size]))
ad_label.append(torch.ones([split_size]))
ad_label = torch.cat(ad_label, dim=0)
ad_label = torch.autograd.Variable(ad_label.cuda(async=True))
a_loss = ad_loss(ad_out, ad_label)
ad_losses.update(a_loss.item(), s_input.size(0)+t_input.size(0))
s_output = output[0:args.batch_size, :]
c_loss = class_loss(s_output, s_target_var)
_, s_pred = torch.max(s_output, dim=1)

# measure accuracy and record loss
prec1 = accuracy(s_pred.data.cpu().long(), s_target.cpu().long(), lmap)
c_losses.update(c_loss.item(), s_input.size(0))
ad_losses.update(a_loss.item(), s_input.size(0))
top1.update(prec1, s_input.size(0))

loss = c_loss + a_loss
# compute gradient and do SGD step
optimizer.zero_grad()

loss.backward()

But since the model definition is in another separate file I’m not sure how to give you an executable code.