@nielsr I hope all is well. May I just follow up on this? What would be the equivalent way of doing that you suggested in pytorch?
tensorflow:
# forward pass
outputs = model(encoding, output_hidden_states=True)

# get the logits
logits = outputs.logits
# get the hidden states
hidden_states = outputs.hidden_states

Thanks!