I managed to introspect the input keys that a model expects by checking model.dummy_inputs
dict (do all models define this reliably? is there a better way?)
But I am currently struggling to find a way to determine the output keys.
The distilbert
model I am experimenting with returns "logits"
, but I understand that not all models may output that key? And also there are config flags (e.g. model.config.output_hidden_states
) that would cause additional keys to be output.
I’m not sure if this is specifically a transformers
question or a PyTorch question…