Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from sentence_transformers import SentenceTransformer, util
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import os
+import torch
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -60,7 +61,7 @@ def generate_response(user_query, relevant_segment):
         input_ids = tokenizer.encode(user_message, return_tensors='pt')
         
         # Create attention mask (1 for real tokens, 0 for padding tokens)
-        attention_mask = (input_ids != tokenizer.pad_token_id).long()
+        attention_mask = (input_ids != tokenizer.pad_token_id).to(torch.long)
 
         # Generate the response using the model
         output = model.generate(