Update README.md

2022-03-10 10:20:37 +01:00
parent fde901877a
commit 6c9010ef63
1 changed files with 2 additions and 1 deletions
--- a/examples/research_projects/jax-projects/README.md
+++ b/examples/research_projects/jax-projects/README.md
@@ -780,7 +780,8 @@ def cross_entropy(logits, labels):
 # define a function which will run the forward pass return loss
 def compute_loss(params, input_ids, labels):
   logits = model(input_ids, params=params, train=True)
-   loss = cross_entropy(logits, onehot(labels)).mean()
+   num_classes = logits.shape[-1]
+   loss = cross_entropy(logits, onehot(labels, num_classes)).mean()
   return loss
 ```