I assume that you are calling model.forward_() in your code instead of the model directly?
If so, then this breaking behavior is expected, as utility methods such as hooks and data parallel rely on the “standard” usage of calling the model directly.
For nn.DataParallel, the model.__call__ method would call into DataParallel.forward, which takes care of the parallelization as seen here 2.