The following worked for me:
def forward(self, X, labels, is_print=False):
        stacked_X = torch.stack(X)
        out = self.transformer(stacked_X)
        with torch.autocast('cuda'):
            loss = self.criterion(out, torch.tensor(labels).cuda())
        pred = out.data.max(1)[1] 
        return pred, labels, loss