Are you sure you’re using both GPUs with nn.DataParallel? The line:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

makes me think that it’s running on a single gpu. You could try changing this to:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

In addition, one of the main benefits of data parallelism is that you can use larger batch sizes to speed up iterating through the dataset, so you can try:
batch_size = 10000 * torch.cuda.device_count()