def before_fit(self): 
        "Set model to cuda before training"                
        self.learner.model.to(self.device)
        self.learner.device = self.device
Set model to cuda before training
PatchTST_self_supervised/src/callback/core.py