How to register a (repeatedly) sampled random tensor?

[RylanSchaeffer]

In my code, with each batch, I sample a random tensor. However, when I try porting to GPU, I get an error that I'm trying to multiply a tensor on CPU with a tensor on GPU. The documentation states

The LightningModule knows what device it is on. You can access the reference via self.device. Sometimes it is necessary to store tensors as module attributes. However, if they are not parameters they will remain on the CPU even if the module gets moved to a new device. To prevent that and remain device agnostic, register the tensor as a buffer in your modules’ init method with register_buffer().

The example given is:

class LitModel(LightningModule):
    def __init__(self):
        ...
        self.register_buffer("sigma", torch.eye(3))
        # you can now access self.sigma anywhere in your module

Using this example, suppose I want to randomly sample a 3x3 matrix sigma with each batch. How do I properly register this tensor?

[RylanSchaeffer]

I just found this answer from 2020: https://stackoverflow.com/questions/63660624/normal-distribution-sampling-in-pytorch-lightning

Is the best solution currently to specify the device?

[akihironitta]

In my code, with each batch, I sample a random tensor. However, when I try porting to GPU, I get an error that I'm trying to multiply a tensor on CPU with a tensor on GPU.

As the thread suggests, you need to move your tensor to GPU manually with self.device in each iteration.

[RylanSchaeffer]

That's the 2020 solution. Is there not a more modern approach?

…

On Tue, Aug 9, 2022, 8:05 PM Akihiro Nitta ***@***.***> wrote: In my code, with each batch, I sample a random tensor. However, when I try porting to GPU, I get an error that I'm trying to multiply a tensor on CPU with a tensor on GPU. As the thread suggests, you need to move your tensor to GPU manually with self.device in each iteration. — Reply to this email directly, view it on GitHub <#14131 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACEHLCYYBOTRQSDD7EPBGPLVYMMBJANCNFSM56C3FX2Q> . You are receiving this because you authored the thread.Message ID: ***@***.*** com>

[akihironitta]

There's no "modern" approach I could think of, and I'd rather say that PyTorch and PyTorch Lightning are stable enough to keep the same APIs for a few years now :)

Or is there any particular feature you'd like to see in PL?

[RylanSchaeffer]

Ok cool thank you! Cheers, Rylan Schaeffer

…

On Tue, Aug 9, 2022 at 8:58 PM Akihiro Nitta ***@***.***> wrote: There's no "modern" approach I could think of, and I'd rather say that PyTorch and PyTorch Lightning are stable enough to keep the same APIs for a few years now :) — Reply to this email directly, view it on GitHub <#14131 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACEHLCZYJOVFMVCBXAKFUILVYMSFJANCNFSM56C3FX2Q> . You are receiving this because you authored the thread.Message ID: ***@***.*** com>