Skip to content

Commit

Permalink
Update shieldgemma_on_huggingface.ipynb
Browse files Browse the repository at this point in the history
fix bug
  • Loading branch information
bebechien authored Aug 1, 2024
1 parent 65e0c0f commit e6f15e5
Showing 1 changed file with 1 addition and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down

0 comments on commit e6f15e5

Please sign in to comment.