Hey @cybaj, the model needs to be moved to GPU before passing it to DataParallel ctor. Have you tried changing the following code:



 cybaj:

self.model = nn.DataParallel(self.model).cuda()


to
self.model = nn.DataParallel(self.model.to("cuda:0"))