Accept `TensorDict` (or more generally, dict-like's) as a `training_step` return type

[egaznep]

Description & Motivation

I have a very similar use case to the one outlined in this issue: #19311 . Different from that, I would like to return a TensorDict (see https://github.com/pytorch/tensordict/) from a training_step. The check inside

pytorch-lightning/src/lightning/pytorch/loops/optimization/manual.py

Line 48 in 37f559e

if isinstance(training_step_output, dict):

requires an dict, but I think it should be a little more flexible to allow a TensorDict.

Pitch

I can imaging having to specify tensordict as a requirement is not preferable, so rather than isinstance(training_step_output, [dict, TensorDict]), I think checking if training_step_output has a keys attribute could work nicely.

Alternatives

No response

Additional context

No response

cc @lantiga @Borda