update to_device

app.py

@@ -152,7 +152,7 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
             }
         ]
         # llama guard check for it
-        # prompt_safety = moderate(chat_round, llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID)['generated_text']
+        prompt_safety = moderate(chat_round, llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID)['generated_text']
         prompt_safety = "safe"
         
         if prompt_safety == "safe":

utils.py

@@ -158,7 +158,7 @@ def load_llama_guard(model_id = "meta-llama/Llama-Guard-3-1B"):
 
     logger.info("loading llama_guard")
     llama_guard_tokenizer = AutoTokenizer.from_pretrained(model_id)
-    llama_guard = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map="cuda")
+    llama_guard = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype).to('cuda')
 
     # Get the id of the "unsafe" token, this will later be used to extract its probability
     UNSAFE_TOKEN_ID = llama_guard_tokenizer.convert_tokens_to_ids("unsafe")