fix: mypy error

kirill-push · Nov 1, 2023 · c1f5eac · c1f5eac
1 parent 7fc9c6b
commit c1f5eac
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/chatbot/model.py b/src/chatbot/model.py
@@ -1,14 +1,15 @@
 from typing import Optional
 
 import torch
+from torch import Tensor
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 class DialoGPTChatbot:
     def __init__(self) -> None:
         self.tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
         self.model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
-        self.chat_history_ids: Optional[torch.Tensor] = None
+        self.chat_history_ids: Optional[Tensor] = None
 
     def get_response(self, user_input: str) -> str:
         new_user_input_ids = self.tokenizer.encode(