How to customize data types that can be transferred between CPU and GPU

ifsheldon: <p>When I design my network I need to work with “complex” data whose leaf attributes are tensors. For example, <code>CameraSpecification</code> contains attributes:</p>
<ul>
<li>position, a tensor</li>
<li>transform, a tensor</li>
<li>intrinsics,  data of another type <code>CameraIntrinsics</code> containing:
<ul>
<li>fov, a tensor</li>
<li>pixel_coordinates, a tensor</li>
</ul>
</li>
</ul>
<p>Then these two classes are nothing but groups of tensors, so <strong>in theory</strong> they can be safely transferred between CPU and GPU and can be generated during <code>forward()</code>. However, if I generate them during <code>forward()</code>, I will have to move the tensor attributes of them using <code>to()</code>. This is especially inconvenient when I use PytorchLightning. Is there a more elegant way to group tensors and to group groups of tensors? Thank you!</p>

ptrblck: <p>I’m not sure, if the posted object is used as the input or inside the model as e.g. a custom parameter.<br>
In the former case, I would guess you could push the data to the GPU before executing the forward pass, while the parameters would be pushed to the device after initializing the model, so I’m unsure why the device transfer is used in the <code>forward</code>.<br>
In any case, I don’t know what limitations Lightning has, so could you explain why pushing the data to the device won’t work?</p>

ifsheldon: <p>I mean, how can I get rid of <code>.to()</code>? Is it a way to tell PyTorch or Lightning that my class is just an aggregation of tensors so then they can just automatically implement and call <code>.to()</code> for me? I was expecting something like this</p>
<pre><code class="lang-python">class TensorGroup(torch.TensorAggregation):
     def __init__(self):
        super().__init__()
        self.tensor1 = torch.rand()
        self.tensor2 = torch.rand()

# in a Lightning training step
def training_step(self,batch, batch_idx):
    some_tensor_group = TensorGroup().to(self.device) # which is the same as some_tensor = torch.rand(1).to(self.device)
    # or even better
    some_tensor_group = TensorGroup() # so I can get rid of .to()
</code></pre>

ptrblck: <aside class="quote no-group" data-username="ifsheldon" data-post="3" data-topic="131069">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ifsheldon/48/41901_2.png" class="avatar"> ifsheldon:</div>
<blockquote>
<p>Is it a way to tell PyTorch or Lightning that my class is just an aggregation of tensors so then they can just automatically implement and call <code>.to()</code> for me?</p>
</blockquote>
</aside>
<p>I don’t know about Lightning, but that’s not possible in PyTorch, but you could create a new custom object and implement the <code>to()</code> method, if you think it’s making the code cleaner somehow.</p>
<p>I still don’t understand the use case, but given your current <code>TensorGroup</code> code snippet, you could try to use a custom <code>nn.Module</code> instead and register the tensors as buffers, which would then push them to the desired device via the <code>to()</code> operation.</p>