Hey,
I’m not sure if this will help or not but i would try the following code:
use_cuda = torch.cuda.device_count() >= 1
device = torch.device("cuda" if use_cuda else "cpu")

netG_A2B = nn.DataParallel(netG_A2B.to(device))
netG_B2A = nn.DataParallel(netG_B2A.to(device))
netD_A = nn.DataParallel(netD_A.to(device))
netD_B = nn.DataParallel(netD_B.to(device))