Overwrite forward in layers to work with Data Parallel

this post showed how to overwrite forward in layers How can I replace the forward method of a predefined torchvision model with my customized forward function? - #6 by Philipp_Friebertshau

hope it helps