I don’t think that’s possible, since the forward hooks are called in the nn.Module.__call__ method before dispatching to the forward method.
Since you are using the functional API you already have the output and input activations and could e.g. store them directly if needed. If hooks are more convenient to use instead of the tensors directly, you would probably have to wrap your calls into a custom module.