Update main.py

main.py

@@ -115,9 +115,9 @@ def getChat(text,model,tokenizer):
     text_ids = tokenizer.convert_tokens_to_ids(text)
     # print(text_ids)
 
-    input_ids = torch.tensor(text_ids).long().to(device)
+    input_ids = torch.tensor(text_ids).long()
     input_ids = input_ids.unsqueeze(0)
-    mask_input = torch.ones_like(input_ids).long().to(device)
+    mask_input = torch.ones_like(input_ids).long()
     # print(input_ids.size())
     response = []  # 根据context，生成的response
     # 最多生成max_len个token