pytorch lightning init not on cuda

1# tensor created in __init__ wont be added to cuda
2# use self.register_buffer to register variables 
3
4class LitModel(LightningModule):
5
6    def __init__(self):
7        ...
8        self.register_buffer("sigma", torch.eye(3))
9        # you can now access self.sigma anywhere in your module