pere
/

norwegian-robert-large-test

Model card Files Files and versions Metrics Training metrics Community

pere commited on 6 days ago

Commit

9868152

•

1 Parent(s): f247126

Update run_mlm_flax.py

Files changed (1) hide show

run_mlm_flax.py +1 -1

run_mlm_flax.py CHANGED Viewed

@@ -796,7 +796,7 @@ def main():
         logits = model(**batch, params=params, train=False)[0]
         # compute loss, ignore padded input tokens
-        label_mask = jnp.where(labels > 0, 1.0, 0.0)
         loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
         # compute accuracy

         logits = model(**batch, params=params, train=False)[0]
         # compute loss, ignore padded input tokens
+        label_mask = jnp.where(labels!=-100, 1.0, 0.0)
         loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
         # compute accuracy