You could use forward hooks as described here which would allow you to check the outputs without changing the forward function in case you are using nn.Modules.
You could use forward hooks as described here which would allow you to check the outputs without changing the forward function in case you are using nn.Modules.