Update app.py
Browse files
app.py
CHANGED
@@ -37,7 +37,8 @@ models = {
|
|
37 |
|
38 |
|
39 |
# Load the mistral model
|
40 |
-
mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
|
|
|
41 |
|
42 |
# Load ESM2 model
|
43 |
model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
|
|
|
37 |
|
38 |
|
39 |
# Load the mistral model
|
40 |
+
mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
|
41 |
+
mistral_model.to('cuda')
|
42 |
|
43 |
# Load ESM2 model
|
44 |
model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
|