Update model.py

model.py

@@ -5,12 +5,16 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+print(f"{torch.cuda.current_device()}")
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
 model_id = "TheBloke/Chronos-Beluga-v2-13B-GPTQ"
 tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
 model = AutoGPTQForCausalLM.from_quantized(model_id,
         use_safetensors=True,
         trust_remote_code=False,
-        device="cuda:0",
+        device="cuda",
         use_triton=False,
         quantize_config=None)