Multi-GPU Tensor Initialization Question
The documentation advises the usage of 'type_as' when initializing new tensors in multi-gpu settings:

https://pytorch-lightning.readthedocs.io/en/stable/advanced/multi_gpu.html#init-tensors-using-type-as-and-register-buffer
When you need to create a new tensor, use type_as. This will make your code scale to any arbitrary number of GPUs or TPUs with Lightning.

The example shows a case where a new tensor is initialized inside a LightningModule forward function:

def forward(self, x):
z = torch.Tensor(2, 3)
z = z.type_as(x)

Presumably x is a tensor that has already been initialized on the target gpu.
My question is what to do in the case where we want to initialize a new tensor on the target gpu, and we do not have access to a tensor that has already been initialized on the target gpu?
For example, how does one properly initialize a new tensor when it is created inside a Dataset constructor that is instantiated during LightningDataModule setup()?

class SomeDataModule(LightningDataModule):
...
def setup(self, stage: Optional[str] = None):
if stage in (None, "fit"):
dataset = SomeDataset()
...

where:

class SomeDataset(Dataset):
def init(self):
self.some_tensor = torch.Tensor(2,3)

Will using type_as on the new tensor initialize the data on the target gpu?

self.some_tensor = self.some_tensor.type_as(self.some_tensor)

Or is a different approach necessary? (e.g. register_buffer())
if it's part of the dataset, it's already moved to the target device when a batch is created while iterating over the dataset.