Fixes torch neural net module for classification tasks only (regression may or may not work properly)