apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-05-03 04:42:49 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix precommit

Browse source
This commit is contained in:
cotran 2024-10-16 16:38:14 -07:00
parent 25fc7a130e
commit 953b1b056b
2 changed files with 7 additions and 103 deletions
Download patch file Download diff file Expand all files Collapse all files

8
model_server/app/tests/test_loaders_gpu.py
View file

@ -5,7 +5,13 @@ import app.commons.globals as glb
from app.loader import get_embedding_model, get_zero_shot_model, get_prompt_guard
# Mock constants
glb.DEVICE = "cuda" # Adjust as needed for your test case
if torch.cuda.is_available():
DEVICE = "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
DEVICE = "mps"
else:
DEVICE = "cpu"
glb.DEVICE = DEVICE # Adjust as needed for your test case
arch_guard_model_type = {
"cpu": "katanemo/Arch-Guard-cpu",
"cuda": "katanemo/Arch-Guard",