pytorch lightning to cuda code example

Example: pytorch lightning init not on cuda

# tensor created in __init__ wont be added to cuda
# use self.register_buffer to register variables 

class LitModel(LightningModule):

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