yuntian-deng commited on
Commit
3ffaf8e
1 Parent(s): 859f68f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -8,6 +8,9 @@ implicit_cot_model_name = 'yuntian-deng/implicit-cot-math-mistral7b'
8
  implicit_cot_model = AutoModelForCausalLM.from_pretrained(implicit_cot_model_name, torch_dtype=torch.bfloat16)
9
  tokenizer = AutoTokenizer.from_pretrained(implicit_cot_model_name)
10
 
 
 
 
11
  # Constants
12
  MAX_RESULT_TOKENS = 10
13
 
 
8
  implicit_cot_model = AutoModelForCausalLM.from_pretrained(implicit_cot_model_name, torch_dtype=torch.bfloat16)
9
  tokenizer = AutoTokenizer.from_pretrained(implicit_cot_model_name)
10
 
11
+ implicit_cot_model.to('cuda' if torch.cuda.is_available() else 'cpu')
12
+ implicit_cot_model.eval()
13
+
14
  # Constants
15
  MAX_RESULT_TOKENS = 10
16