diff --git a/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb b/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb
index dbfa3e911..904b9a955 100644
--- a/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb
+++ b/site/en/responsible/docs/safeguards/shieldgemma_on_huggingface.ipynb
@@ -236,7 +236,7 @@
         "\n",
         "  # Get logits. Shape <float>[batch_size, sequnece_length, vocab_size]\n",
         "  with torch.no_grad():\n",
-        "    logits = model(**inputs).logits\n",
+        "    logits = shieldgemma(**inputs).logits\n",
         "\n",
         "  # Extract the logits for the Yes and No tokens\n",
         "  yes_no_logits = logits[0, -1, [YES_TOKEN_IDX, NO_TOKEN_IDX]]\n",