Great! There are 3 popular errors on that code block:
- Reversing the arguments.
- Forgetting the transposes.
- Forgetting the from_logits = True argument.
Or some combination of the above. Each mistake gives a characteristic wrong answer, but the combinations can be a little harder to diagnose.