`torch.nn.DataParallel` is deprecated and doesn't work well in many systems with newer GPUs. Get rid of this and find alternative solutions.