You could use forward hooks as described in this post or rename the output activation and return it in the forward method. The latter approach might be useful in case you would like to always return this activation, while the former one is more ad hoc.