Saving weights of epoch 1 at step 92

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fd2e92a7233e29a509eb154637ad78cef153ebf1065ed18cb43fa82960412b88
+oid sha256:51d86bd352715e1623b69a8451f8c752c314bb6cf7669a5d9bb2f7589261d8c3
 size 1419367919

results_tensorboard/events.out.tfevents.1626415968.t1v-n-8cb15980-w-0.843954.3.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7ba3313ae634173f24d2dd23ebaba16f5e2889922e2e54aa323a38a726b39e1
+size 40

results_tensorboard/events.out.tfevents.1626416209.t1v-n-8cb15980-w-0.845549.3.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10501817e5f6a56a85f8dbb3b994e577cc9eb1629d28bb96dc8ad403f11b4dda
+size 25152

src/prediction.py

@@ -1,4 +1,5 @@
 import jax
+print(jax.local_device_count())
 import jax.numpy as jnp
 
 import flax

src/test.py

@@ -1,4 +1,5 @@
 import jax
+print(jax.local_device_count())
 import jax.numpy as jnp
 
 import flax

src/train.py

@@ -111,7 +111,7 @@ def main():
   def eval_function(logits):
     return logits.argmax(-1)
 
-  model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2-large',input_shape=(1,4,1))
+  model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2-medium',input_shape=(1,4,1))
 
   state=TrainState.create(apply_fn=model.__call__,
                           params=model.params,
@@ -238,4 +238,4 @@ def main():
   summary_writer.flush()
 
 if __name__ == "__main__":
-    main()
+    main()