Update app.py

app.py

@@ -66,7 +66,7 @@ class ProteinDataset(Dataset):
         
         # Create labels (tokens for ground truth AAs)
         label_seq = protein_seq + peptide_seq
-        labels = self.tokenizer(label_seq, return_tensors="pt", padding="max_length", max_length = self.max_length_pm, truncation=True)["input_ids"].squeeze()
+        labels = self.tokenizer(label_seq, return_tensors="pt", padding="max_length", max_length = self.max_length_pm, truncation=True)["input_ids"].to(device).squeeze()
         
         # Set non-masked positions in the labels tensor to -100
         labels = torch.where(input_ids == self.tokenizer.mask_token_id, labels, -100)