How can create a data in the nn.model object and move it to gpu?

implementation help
1

how can create a data in the nn.model object and move it to gpu?

torch.arange(seq_len)
how can move it GPU, while this is inside my forward method?

2

A better way to handle this would be to use this:

torch.arrange(seq_len, device=self.device)

LightningModule has a property device which can be accessed using self.device which would be of use, using this would make your module device-agnostic at the same time you can make full use of the functionality provided by Lightning.

2 Likes