Update app.py

app.py

@@ -16,11 +16,12 @@ bnb_config = BitsAndBytesConfig(
 
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 # Load model in this way if use GPU
-if (device == "cuda"):
+if torch.cuda.is_available():
     model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=bnb_config)
-elif(device == "cpu"):
+else:
     model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True)
 # Move the model to the GPU if available
+model.to(device)
 
 generation_config = GenerationConfig(
     penalty_alpha=0.6,