Update app.py

app.py

@@ -16,6 +16,8 @@ explicit_cot_model = AutoModelForCausalLM.from_pretrained(explicit_cot_model_nam
 
 models = {'implicit': implicit_cot_model, 'no': no_cot_model, 'explicit': explicit_cot_model}
 
+[model.to('cuda' if torch.cuda.is_available() else 'cpu') for model in models.values()]
+
 # Constants
 MAX_PRODUCT_DIGITS_PER_MODEL = {'implicit': 100, 'no': 100, 'explicit': 960}