Update services/llama_generator.py

services/llama_generator.py

@@ -4,6 +4,8 @@ from config.config import GenerationConfig, ModelConfig
 from typing import List, Dict, Any, Optional, Tuple
 from datetime import datetime
 import logging
+import pytorch
+
 from config.config import settings
 
 from services.prompt_builder import LlamaPromptTemplate
@@ -50,8 +52,9 @@ class LlamaGenerator(BaseGenerator):
         print(llama_model_name)
         print(prm_model_path)
 
-        self.model_manager = ModelManager()
-            
+        self.model_manager = ModelManager() 
+        
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.tokenizer = self.model_manager.load_tokenizer(llama_model_name)  # Add this line to initialize the tokenizer
  
         super().__init__(