Saving weights and logs of step 5000

events.out.tfevents.1734082062.t1v-n-53cd541d-w-35.1086346.0.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:22d24708144113ab237002eb807b1d5c7a911ce761f98ce2daadc5f2d1a7c3ed
-size 63038
+oid sha256:b85b9a72fe6b897c2e4406605020d7c710cb9eebc969bde09e273fb974300238
+size 228898

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3f1adcfdad8a681e08bfe5535021067e1c876db067eda92f40f92ccf9c4c5e63
+oid sha256:8c4c381d4203e6b03a768f5cd7066942ed7e02f54c5534d09efd039be2c86d2e
 size 1421658229

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