Spaces:

ibm-granite
/

granite-guardian-3.1-8b

Running on Zero

Martín Santillán Cooper commited on 19 days ago

Commit

809aedc

•

1 Parent(s): 0f55f50

Change cpu -> cuda

Files changed (1) hide show

src/model.py CHANGED Viewed

@@ -19,7 +19,7 @@ logger.debug(f"Inference engine is: '{inference_engine}'")
 if inference_engine == "TORCH":
     import torch
-    device = torch.device("cpu")
     model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.1-8b")
     logger.debug(f"model_path is {model_path}")

 if inference_engine == "TORCH":
     import torch
+    device = torch.device("cuda")
     model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.1-8b")
     logger.debug(f"model_path is {model_path}")