Thanks for your response ptrblck!
This is actually part of the inference, and not training.
Originally the entire model was in a .tar format and was being loaded using torch.nn.DataParallel. But since we only had 1 GPU and are trying to cut down as much time as possible, I re saved the checkpoint only using the following:
checkpoint = torch.load(self.validate_model)
self.model.module.load_state_dict(checkpoint[‘state_dict’])
torch.save(self.model.module.state_dict(), path)
And then loading this newly saved state dict for inference without using torch.nn.DataParallel.
Is this not the right way to approach this? I am new to PyTorch.
Thanks!