Update config.py

config.py

@@ -3,7 +3,7 @@
 import torch
 
 # Device configuration
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Path to the Yi-Coder model
 MODEL_PATH = "01-ai/Yi-Coder-9B-Chat"  # Replace with your actual model path
@@ -12,4 +12,4 @@ MODEL_PATH = "01-ai/Yi-Coder-9B-Chat"  # Replace with your actual model path
 MAX_LENGTH = 1024
 
 # Temperature for text generation
-TEMPERATURE = 0.7
+TEMPERATURE = 0.7