Add Support for Dataclasses to Trainer

[ntenenz]

Feature request

Update Trainer._prepare_input to natively support python dataclasses for support of more structured objects.

Motivation

_prepare_input will seamlessly transfer the tensors contained in many datatypes (list, tuple, dict, etc.) to the appropriate device. However, it will not do so for dataclass objects.

Python dataclasses often provide better ergonomics than TypedDict, which is likely the closest supported alternative. Adding support appears to be a small change to the codebase with nice benefits to the user.

Your contribution

If the PR is as simple as it initially appears, I would be willing to submit a PR.