This might work, but note that it’s not a general rule since your predict method is a custom function which is not a standard method in nn.Module.
The usual and required workflow is to override the forward method. You should thus compare your predict method with the forward and check what the difference is.
nn.DataParallel doesn’t have any knowledge about custom methods (such as predict) and thus will use the standard forward method to use the data parallel approach.