Update run_mlm_flax.py

run_mlm_flax.py

@@ -760,7 +760,7 @@ def main():
             logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
 
             # compute loss, ignore padded input tokens
-            label_mask = jnp.where(labels > -100, 1.0, 0.0)
+            label_mask = jnp.where(labels!=-100, 1.0, 0.0)
             loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
 
             # take average