See here and here
basically you have to put your model into a dataparallel temporarily again or change the key names to remove “module.”