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:7566bc459faaabe05b6f642a97a367b444365761679f1b5dd13312d70b413601
+oid sha256:51d86bd352715e1623b69a8451f8c752c314bb6cf7669a5d9bb2f7589261d8c3
 size 1419367919

model_file.py

@@ -190,7 +190,7 @@ class FlaxGPT2ForMultipleChoiceModule(nn.Module):
   dtype: jnp.dtype = jnp.float32
   def setup(self):
     self.transformer = FlaxGPT2Module(config=self.config, dtype=self.dtype)
-    self.dropout = nn.Dropout(rate=0.3)
+    self.dropout = nn.Dropout(rate=0.2)
     self.classifier = nn.Dense(4, dtype=self.dtype)
 
   def __call__(self,input_ids,attention_mask,position_ids,return_dict=True,deterministic=True,*args):

results_tensorboard/events.out.tfevents.1626339960.t1v-n-8cb15980-w-0.776261.3.v2

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

train.py

@@ -74,7 +74,7 @@ def main():
   per_device_batch_size=4
   seed=0
   num_train_epochs=3
-  learning_rate=4e-5
+  learning_rate=2e-5
   
 
   total_batch_size = per_device_batch_size * jax.local_device_count()