Keep the data structure of tensor-like objects after moving them to devices
🚀 Feature
For now, PL handles the device conversion of data pretty good and make it almost transparent to users. But there are still restrictions like 1. data must be instances of torch.Tensor (or its subclasses) to be convertible, 2. after conversion, the data lose the original structures (and become torch.Tensor even if they are actually instances of its subclasses).
So I wonder if we can relax the limitation of the data, for example, to those data structures that have ".to()" method, and thus can preserve the original data structures after device conversion.
Motivation
Make it easier to utilize structured data, like those defined in detectron2.
Pitch
Keep the original class information after the data conversion
Alternatives
If the relaxation is too loose, it's okay to keep the type check of torch.Tensor