Update app.py

app.py

@@ -1,6 +1,9 @@
 import torch
 from transformers import TextIteratorStreamer
 import threading
+import spaces
+
+print("Is CUDA available?", torch.cuda.is_available())
 
 class ModelWrapper:
     def __init__(self):
@@ -9,16 +12,18 @@ class ModelWrapper:
     @spaces.GPU
     def generate(self, prompt):
         if self.model is None:
-            # Load the model when GPU is allocated
+            # Explicitly set device_map to 'cuda'
             self.model = AutoGPTQForCausalLM.from_quantized(
                 model_id,
-                device_map='auto',
+                device_map={'': 'cuda:0'},
                 trust_remote_code=True,
             )
-            self.model.eval()
+
+        print("Model is on device:", next(self.model.parameters()).device)
 
         # Tokenize the input prompt
         inputs = tokenizer(prompt, return_tensors='pt').to('cuda')
+        print("Inputs are on device:", inputs['input_ids'].device)
 
         # Set up the streamer
         streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)