Access to intermediate embeddings

[blazejba]

This change allows access to intermediate embeddings from the model. Example of usage:

output, extras = model(batch, extra_return_names=["pre_task_heads"])
fingerprint_graph = extras['pre_task_heads']['graph_feat']

[DomInvivo]

It doesn't seem to break anything, so fine by me, but @WenkelF would likely be better suited to approve it.

[WenkelF]

I agree with @DomInvivo that the changes don't break anything (although I see the unit tests failing).

Out of curiosity, what is the purpose of the changes @blazejba? It looks like they make the input of the task heads accessible as fingerprints when running a forward pass.

We also have this logic for fingerprinting. It allows not only to read out the inputs of the task heads but any intermediate representation of the model. Maybe there is a way of consolidating here if it serves the same purpose to avoid redundancies. But I am also ok with the changes as is.

[blazejba]

@WenkelF - Thanks for the pointer. I wasn't aware it's already implemented - it's exactly what I've been looking for. I am closing this PR as it is no longer necessary.