Skip to content

Commit a654467

Browse files
authored
feat: add cpu/cuda config for prompt guard (#2194)
# What does this PR do? Previously prompt guard was hard coded to require cuda which prevented it from being used on an instance without a cuda support. This PR allows prompt guard to be configured to use either cpu or cuda. [//]: # (If resolving an issue, uncomment and update the line below) Closes [#2133](#2133) ## Test Plan (Edited after incorporating suggestion) 1) started stack configured with prompt guard as follows on a system without a GPU and validated prompt guard could be used through the APIs 2) validated on a system with a gpu (but without llama stack) that the python selecting between cpu and cuda support returned the right value when a cuda device was available. 3) ran the unit tests as per - https://github.com/meta-llama/llama-stack/blob/main/tests/unit/README.md [//]: # (## Documentation) --------- Signed-off-by: Michael Dawson <[email protected]>
1 parent 63a9f08 commit a654467

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,9 @@ def __init__(
7575
self.temperature = temperature
7676
self.threshold = threshold
7777

78-
self.device = "cuda"
78+
self.device = "cpu"
79+
if torch.cuda.is_available():
80+
self.device = "cuda"
7981

8082
# load model and tokenizer
8183
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)

0 commit comments

Comments
 (0)