I’m a bit confused, what these lines are doing:
wrapped_model = nn.DataParallel(learner.model)
learner.model = wrapped_model.module

I assume that learner.model is an nn.Module, wrapped_model would be an nn.DataParallel object, and you would reassign the module to itself?
What would happen, if you remove the second line of code?