diff --git "a/trainer_state.json" "b/trainer_state.json"
--- "a/trainer_state.json"
+++ "b/trainer_state.json"
@@ -11,3873 +11,3873 @@
   "log_history": [
     {
       "epoch": 0.0018066847335140017,
-      "grad_norm": 2.6597820350894357,
+      "grad_norm": 2.906479789331899,
       "learning_rate": 0.0,
       "loss": 0.8251,
       "step": 1
     },
     {
       "epoch": 0.0036133694670280035,
-      "grad_norm": 2.7266510344811494,
+      "grad_norm": 2.946569058957928,
       "learning_rate": 1.0714285714285714e-06,
       "loss": 0.8284,
       "step": 2
     },
     {
       "epoch": 0.005420054200542005,
-      "grad_norm": 2.7200407943823968,
+      "grad_norm": 2.96702642089792,
       "learning_rate": 2.1428571428571427e-06,
-      "loss": 0.8429,
+      "loss": 0.8428,
       "step": 3
     },
     {
       "epoch": 0.007226738934056007,
-      "grad_norm": 2.6398267415203076,
+      "grad_norm": 2.9084398756854015,
       "learning_rate": 3.2142857142857143e-06,
-      "loss": 0.8381,
+      "loss": 0.8378,
       "step": 4
     },
     {
       "epoch": 0.009033423667570008,
-      "grad_norm": 2.450058379558589,
+      "grad_norm": 2.621568341450903,
       "learning_rate": 4.2857142857142855e-06,
-      "loss": 0.8245,
+      "loss": 0.8229,
       "step": 5
     },
     {
       "epoch": 0.01084010840108401,
-      "grad_norm": 1.9536840033460428,
+      "grad_norm": 2.1930464956581446,
       "learning_rate": 5.357142857142857e-06,
-      "loss": 0.814,
+      "loss": 0.8097,
       "step": 6
     },
     {
       "epoch": 0.012646793134598013,
-      "grad_norm": 1.2180340367347295,
+      "grad_norm": 1.4933984713931483,
       "learning_rate": 6.428571428571429e-06,
-      "loss": 0.7845,
+      "loss": 0.7742,
       "step": 7
     },
     {
       "epoch": 0.014453477868112014,
-      "grad_norm": 1.1514973199014618,
+      "grad_norm": 1.397985264834512,
       "learning_rate": 7.5e-06,
-      "loss": 0.7874,
+      "loss": 0.7739,
       "step": 8
     },
     {
       "epoch": 0.016260162601626018,
-      "grad_norm": 1.1427251731602237,
+      "grad_norm": 1.3545488436805673,
       "learning_rate": 8.571428571428571e-06,
-      "loss": 0.7793,
+      "loss": 0.7663,
       "step": 9
     },
     {
       "epoch": 0.018066847335140017,
-      "grad_norm": 2.4537001853666673,
+      "grad_norm": 2.316903712014158,
       "learning_rate": 9.642857142857144e-06,
-      "loss": 0.7624,
+      "loss": 0.7505,
       "step": 10
     },
     {
       "epoch": 0.01987353206865402,
-      "grad_norm": 2.474412722550665,
+      "grad_norm": 2.3724937294713033,
       "learning_rate": 1.0714285714285714e-05,
-      "loss": 0.7491,
+      "loss": 0.7368,
       "step": 11
     },
     {
       "epoch": 0.02168021680216802,
-      "grad_norm": 2.333456673377541,
+      "grad_norm": 2.1948974279154383,
       "learning_rate": 1.1785714285714286e-05,
-      "loss": 0.762,
+      "loss": 0.7501,
       "step": 12
     },
     {
       "epoch": 0.023486901535682024,
-      "grad_norm": 1.6274020508818596,
+      "grad_norm": 1.7769826234474873,
       "learning_rate": 1.2857142857142857e-05,
-      "loss": 0.7368,
+      "loss": 0.7263,
       "step": 13
     },
     {
       "epoch": 0.025293586269196026,
-      "grad_norm": 1.563869927563279,
+      "grad_norm": 1.7537734276506007,
       "learning_rate": 1.3928571428571429e-05,
-      "loss": 0.7378,
+      "loss": 0.7297,
       "step": 14
     },
     {
       "epoch": 0.02710027100271003,
-      "grad_norm": 1.253626680340628,
+      "grad_norm": 1.4150131497465723,
       "learning_rate": 1.5e-05,
-      "loss": 0.7271,
+      "loss": 0.7181,
       "step": 15
     },
     {
       "epoch": 0.028906955736224028,
-      "grad_norm": 0.9176330164857595,
+      "grad_norm": 1.388901011985526,
       "learning_rate": 1.6071428571428572e-05,
-      "loss": 0.7131,
+      "loss": 0.7029,
       "step": 16
     },
     {
       "epoch": 0.03071364046973803,
-      "grad_norm": 0.7185343991261351,
+      "grad_norm": 0.9533064996374424,
       "learning_rate": 1.7142857142857142e-05,
-      "loss": 0.7131,
+      "loss": 0.7021,
       "step": 17
     },
     {
       "epoch": 0.032520325203252036,
-      "grad_norm": 0.7884609645105337,
+      "grad_norm": 0.7772119179800466,
       "learning_rate": 1.8214285714285712e-05,
-      "loss": 0.7045,
+      "loss": 0.6907,
       "step": 18
     },
     {
       "epoch": 0.03432700993676603,
-      "grad_norm": 0.7399830971167167,
+      "grad_norm": 0.7719134453183193,
       "learning_rate": 1.928571428571429e-05,
-      "loss": 0.6995,
+      "loss": 0.6849,
       "step": 19
     },
     {
       "epoch": 0.036133694670280034,
-      "grad_norm": 0.7376328500886337,
+      "grad_norm": 0.6865260996169726,
       "learning_rate": 2.0357142857142858e-05,
-      "loss": 0.6779,
+      "loss": 0.661,
       "step": 20
     },
     {
       "epoch": 0.037940379403794036,
-      "grad_norm": 0.6439873145292982,
+      "grad_norm": 0.6261375647490827,
       "learning_rate": 2.1428571428571428e-05,
-      "loss": 0.6834,
+      "loss": 0.6666,
       "step": 21
     },
     {
       "epoch": 0.03974706413730804,
-      "grad_norm": 0.5790542088592184,
+      "grad_norm": 0.6378186989443875,
       "learning_rate": 2.25e-05,
-      "loss": 0.6721,
+      "loss": 0.6558,
       "step": 22
     },
     {
       "epoch": 0.04155374887082204,
-      "grad_norm": 0.5191304250088061,
+      "grad_norm": 0.6713458173721598,
       "learning_rate": 2.357142857142857e-05,
-      "loss": 0.6707,
+      "loss": 0.6564,
       "step": 23
     },
     {
       "epoch": 0.04336043360433604,
-      "grad_norm": 0.4488946670945292,
+      "grad_norm": 0.6274625379226225,
       "learning_rate": 2.464285714285714e-05,
-      "loss": 0.676,
+      "loss": 0.6614,
       "step": 24
     },
     {
       "epoch": 0.045167118337850046,
-      "grad_norm": 0.4798647764154148,
+      "grad_norm": 0.5715761838454848,
       "learning_rate": 2.5714285714285714e-05,
-      "loss": 0.6585,
+      "loss": 0.6432,
       "step": 25
     },
     {
       "epoch": 0.04697380307136405,
-      "grad_norm": 0.4204330566493729,
+      "grad_norm": 0.48739002531211656,
       "learning_rate": 2.6785714285714288e-05,
-      "loss": 0.673,
+      "loss": 0.6574,
       "step": 26
     },
     {
       "epoch": 0.04878048780487805,
-      "grad_norm": 0.47130112684574876,
+      "grad_norm": 0.5170424711902095,
       "learning_rate": 2.7857142857142858e-05,
-      "loss": 0.6516,
+      "loss": 0.6355,
       "step": 27
     },
     {
       "epoch": 0.05058717253839205,
-      "grad_norm": 0.4362815754533529,
+      "grad_norm": 0.48953368197603536,
       "learning_rate": 2.892857142857143e-05,
-      "loss": 0.6545,
+      "loss": 0.6386,
       "step": 28
     },
     {
       "epoch": 0.052393857271906055,
-      "grad_norm": 0.43822648560890193,
+      "grad_norm": 0.5005989792055041,
       "learning_rate": 3e-05,
-      "loss": 0.6647,
+      "loss": 0.6504,
       "step": 29
     },
     {
       "epoch": 0.05420054200542006,
-      "grad_norm": 0.3853311574389301,
+      "grad_norm": 0.4544942981478791,
       "learning_rate": 2.9999731440137413e-05,
-      "loss": 0.6644,
+      "loss": 0.6514,
       "step": 30
     },
     {
       "epoch": 0.05600722673893405,
-      "grad_norm": 0.4221785516753897,
+      "grad_norm": 0.3974914740221412,
       "learning_rate": 2.9998925770166232e-05,
-      "loss": 0.6534,
+      "loss": 0.6402,
       "step": 31
     },
     {
       "epoch": 0.057813911472448055,
-      "grad_norm": 0.3980570059327311,
+      "grad_norm": 0.3878954325744685,
       "learning_rate": 2.9997583018935875e-05,
-      "loss": 0.6416,
+      "loss": 0.6281,
       "step": 32
     },
     {
       "epoch": 0.05962059620596206,
-      "grad_norm": 0.35571676272289743,
+      "grad_norm": 0.4561986805367028,
       "learning_rate": 2.9995703234527553e-05,
-      "loss": 0.6513,
+      "loss": 0.6383,
       "step": 33
     },
     {
       "epoch": 0.06142728093947606,
-      "grad_norm": 0.38024957472516396,
+      "grad_norm": 0.37737783307889894,
       "learning_rate": 2.999328648425255e-05,
-      "loss": 0.6385,
+      "loss": 0.6262,
       "step": 34
     },
     {
       "epoch": 0.06323396567299007,
-      "grad_norm": 0.2924858728569485,
+      "grad_norm": 0.3869373298490812,
       "learning_rate": 2.999033285464982e-05,
-      "loss": 0.644,
+      "loss": 0.6313,
       "step": 35
     },
     {
       "epoch": 0.06504065040650407,
-      "grad_norm": 0.3482798299572746,
+      "grad_norm": 0.4272515794130934,
       "learning_rate": 2.9986842451482876e-05,
-      "loss": 0.6387,
+      "loss": 0.6274,
       "step": 36
     },
     {
       "epoch": 0.06684733514001806,
-      "grad_norm": 0.3705337140725527,
+      "grad_norm": 0.36158577191319774,
       "learning_rate": 2.9982815399736008e-05,
-      "loss": 0.6432,
+      "loss": 0.632,
       "step": 37
     },
     {
       "epoch": 0.06865401987353206,
-      "grad_norm": 0.3364721163164359,
+      "grad_norm": 0.4217112984330933,
       "learning_rate": 2.9978251843609816e-05,
-      "loss": 0.6351,
+      "loss": 0.6248,
       "step": 38
     },
     {
       "epoch": 0.07046070460704607,
-      "grad_norm": 0.3272631660096153,
+      "grad_norm": 0.7204159027715281,
       "learning_rate": 2.9973151946516027e-05,
-      "loss": 0.6364,
+      "loss": 0.6269,
       "step": 39
     },
     {
       "epoch": 0.07226738934056007,
-      "grad_norm": 0.29708127282462954,
+      "grad_norm": 0.3648244667440142,
       "learning_rate": 2.996751589107167e-05,
-      "loss": 0.6252,
+      "loss": 0.6154,
       "step": 40
     },
     {
       "epoch": 0.07407407407407407,
-      "grad_norm": 0.24921077158970448,
+      "grad_norm": 0.37700089354074445,
       "learning_rate": 2.9961343879092512e-05,
-      "loss": 0.6308,
+      "loss": 0.6218,
       "step": 41
     },
     {
       "epoch": 0.07588075880758807,
-      "grad_norm": 0.3479569257927118,
+      "grad_norm": 0.32409600863650384,
       "learning_rate": 2.9954636131585845e-05,
-      "loss": 0.6234,
+      "loss": 0.6143,
       "step": 42
     },
     {
       "epoch": 0.07768744354110207,
-      "grad_norm": 0.28917511643628596,
+      "grad_norm": 0.3707207340219278,
       "learning_rate": 2.9947392888742566e-05,
-      "loss": 0.6233,
+      "loss": 0.6148,
       "step": 43
     },
     {
       "epoch": 0.07949412827461608,
-      "grad_norm": 0.34737822266968965,
+      "grad_norm": 0.39427784806992083,
       "learning_rate": 2.993961440992859e-05,
-      "loss": 0.6171,
+      "loss": 0.609,
       "step": 44
     },
     {
       "epoch": 0.08130081300813008,
-      "grad_norm": 0.3247689228102166,
+      "grad_norm": 0.32610161762849305,
       "learning_rate": 2.993130097367553e-05,
-      "loss": 0.6238,
+      "loss": 0.6168,
       "step": 45
     },
     {
       "epoch": 0.08310749774164408,
-      "grad_norm": 0.30719902123424325,
+      "grad_norm": 0.35669866117775423,
       "learning_rate": 2.9922452877670775e-05,
-      "loss": 0.6252,
+      "loss": 0.6185,
       "step": 46
     },
     {
       "epoch": 0.08491418247515808,
-      "grad_norm": 0.3175141153740525,
+      "grad_norm": 0.3676448513773045,
       "learning_rate": 2.991307043874677e-05,
-      "loss": 0.6281,
+      "loss": 0.6213,
       "step": 47
     },
     {
       "epoch": 0.08672086720867209,
-      "grad_norm": 0.31244488771705314,
+      "grad_norm": 0.38768747406718834,
       "learning_rate": 2.9903153992869734e-05,
-      "loss": 0.6148,
+      "loss": 0.6072,
       "step": 48
     },
     {
       "epoch": 0.08852755194218609,
-      "grad_norm": 0.3429131319420423,
+      "grad_norm": 0.3564630767927543,
       "learning_rate": 2.989270389512756e-05,
-      "loss": 0.6211,
+      "loss": 0.6154,
       "step": 49
     },
     {
       "epoch": 0.09033423667570009,
-      "grad_norm": 0.3104062918234721,
+      "grad_norm": 0.38941497566142114,
       "learning_rate": 2.988172051971717e-05,
-      "loss": 0.6204,
+      "loss": 0.615,
       "step": 50
     },
     {
       "epoch": 0.0921409214092141,
-      "grad_norm": 0.31461103786584654,
+      "grad_norm": 0.3584473583813607,
       "learning_rate": 2.9870204259931062e-05,
-      "loss": 0.6097,
+      "loss": 0.6041,
       "step": 51
     },
     {
       "epoch": 0.0939476061427281,
-      "grad_norm": 0.3151625720082357,
+      "grad_norm": 0.6346955736685954,
       "learning_rate": 2.9858155528143256e-05,
-      "loss": 0.6268,
+      "loss": 0.6216,
       "step": 52
     },
     {
       "epoch": 0.0957542908762421,
-      "grad_norm": 0.2860126481493585,
+      "grad_norm": 0.41395691562612613,
       "learning_rate": 2.9845574755794522e-05,
-      "loss": 0.6125,
+      "loss": 0.6074,
       "step": 53
     },
     {
       "epoch": 0.0975609756097561,
-      "grad_norm": 0.28841121592431124,
+      "grad_norm": 0.3325284883699586,
       "learning_rate": 2.9832462393376926e-05,
-      "loss": 0.6018,
+      "loss": 0.5977,
       "step": 54
     },
     {
       "epoch": 0.0993676603432701,
-      "grad_norm": 0.23534181927847642,
+      "grad_norm": 0.4006744277449496,
       "learning_rate": 2.9818818910417706e-05,
-      "loss": 0.6169,
+      "loss": 0.6134,
       "step": 55
     },
     {
       "epoch": 0.1011743450767841,
-      "grad_norm": 0.28661471046317555,
+      "grad_norm": 0.3503235962002047,
       "learning_rate": 2.9804644795462437e-05,
-      "loss": 0.6212,
+      "loss": 0.6178,
       "step": 56
     },
     {
       "epoch": 0.10298102981029811,
-      "grad_norm": 0.2443857675436309,
+      "grad_norm": 0.3596793230653787,
       "learning_rate": 2.9789940556057574e-05,
-      "loss": 0.6065,
+      "loss": 0.6026,
       "step": 57
     },
     {
       "epoch": 0.10478771454381211,
-      "grad_norm": 0.292418070185767,
+      "grad_norm": 0.375719190862958,
       "learning_rate": 2.9774706718732255e-05,
-      "loss": 0.6152,
+      "loss": 0.6115,
       "step": 58
     },
     {
       "epoch": 0.10659439927732611,
-      "grad_norm": 0.24742182455944323,
+      "grad_norm": 0.40865738147424024,
       "learning_rate": 2.9758943828979444e-05,
-      "loss": 0.6009,
+      "loss": 0.5972,
       "step": 59
     },
     {
       "epoch": 0.10840108401084012,
-      "grad_norm": 0.25007211477600044,
+      "grad_norm": 0.459336272208224,
       "learning_rate": 2.9742652451236414e-05,
-      "loss": 0.6023,
+      "loss": 0.5984,
       "step": 60
     },
     {
       "epoch": 0.1102077687443541,
-      "grad_norm": 0.24539117288225054,
+      "grad_norm": 0.4395461744005695,
       "learning_rate": 2.972583316886451e-05,
-      "loss": 0.6059,
+      "loss": 0.6027,
       "step": 61
     },
     {
       "epoch": 0.1120144534778681,
-      "grad_norm": 0.2544907417129908,
+      "grad_norm": 0.5409276356046155,
       "learning_rate": 2.9708486584128303e-05,
-      "loss": 0.6107,
+      "loss": 0.607,
       "step": 62
     },
     {
       "epoch": 0.11382113821138211,
-      "grad_norm": 0.25269030392560393,
+      "grad_norm": 0.4127832467579958,
       "learning_rate": 2.9690613318173966e-05,
-      "loss": 0.611,
+      "loss": 0.6072,
       "step": 63
     },
     {
       "epoch": 0.11562782294489611,
-      "grad_norm": 0.20651691272156247,
+      "grad_norm": 0.3295124749236948,
       "learning_rate": 2.9672214011007087e-05,
-      "loss": 0.6091,
+      "loss": 0.606,
       "step": 64
     },
     {
       "epoch": 0.11743450767841011,
-      "grad_norm": 0.25759031036445135,
+      "grad_norm": 0.38489242897645604,
       "learning_rate": 2.9653289321469715e-05,
-      "loss": 0.6031,
+      "loss": 0.5998,
       "step": 65
     },
     {
       "epoch": 0.11924119241192412,
-      "grad_norm": 0.27004948462624506,
+      "grad_norm": 0.35627109239374033,
       "learning_rate": 2.9633839927216793e-05,
-      "loss": 0.6061,
+      "loss": 0.6026,
       "step": 66
     },
     {
       "epoch": 0.12104787714543812,
-      "grad_norm": 0.23899477664304233,
+      "grad_norm": 0.2956644369843949,
       "learning_rate": 2.9613866524691867e-05,
-      "loss": 0.5958,
+      "loss": 0.593,
       "step": 67
     },
     {
       "epoch": 0.12285456187895212,
-      "grad_norm": 0.3109356502387666,
+      "grad_norm": 0.3766640456122109,
       "learning_rate": 2.9593369829102173e-05,
-      "loss": 0.6119,
+      "loss": 0.6093,
       "step": 68
     },
     {
       "epoch": 0.12466124661246612,
-      "grad_norm": 0.23973937288759506,
+      "grad_norm": 0.3452101448577207,
       "learning_rate": 2.957235057439301e-05,
-      "loss": 0.6135,
+      "loss": 0.6105,
       "step": 69
     },
     {
       "epoch": 0.12646793134598014,
-      "grad_norm": 0.29402451520285133,
+      "grad_norm": 0.3069425861118418,
       "learning_rate": 2.955080951322147e-05,
-      "loss": 0.6178,
+      "loss": 0.6148,
       "step": 70
     },
     {
       "epoch": 0.12827461607949414,
-      "grad_norm": 0.29085498534508775,
+      "grad_norm": 0.41522461019104656,
       "learning_rate": 2.9528747416929467e-05,
-      "loss": 0.5907,
+      "loss": 0.5875,
       "step": 71
     },
     {
       "epoch": 0.13008130081300814,
-      "grad_norm": 0.25477514255751904,
+      "grad_norm": 0.35892642901243405,
       "learning_rate": 2.9506165075516148e-05,
-      "loss": 0.6073,
+      "loss": 0.6043,
       "step": 72
     },
     {
       "epoch": 0.13188798554652212,
-      "grad_norm": 0.28366229733780235,
+      "grad_norm": 0.3560421570359481,
       "learning_rate": 2.9483063297609577e-05,
-      "loss": 0.6063,
+      "loss": 0.6034,
       "step": 73
     },
     {
       "epoch": 0.13369467028003612,
-      "grad_norm": 0.26317649422081363,
+      "grad_norm": 0.3724223713554414,
       "learning_rate": 2.9459442910437798e-05,
-      "loss": 0.598,
+      "loss": 0.5958,
       "step": 74
     },
     {
       "epoch": 0.13550135501355012,
-      "grad_norm": 0.28039599919104496,
+      "grad_norm": 0.33938996629688567,
       "learning_rate": 2.94353047597992e-05,
-      "loss": 0.5971,
+      "loss": 0.5944,
       "step": 75
     },
     {
       "epoch": 0.13730803974706413,
-      "grad_norm": 0.271745198157929,
+      "grad_norm": 0.3499978376673935,
       "learning_rate": 2.941064971003224e-05,
-      "loss": 0.593,
+      "loss": 0.5908,
       "step": 76
     },
     {
       "epoch": 0.13911472448057813,
-      "grad_norm": 0.24998502389958926,
+      "grad_norm": 0.3727209187041342,
       "learning_rate": 2.9385478643984484e-05,
-      "loss": 0.5878,
+      "loss": 0.5866,
       "step": 77
     },
     {
       "epoch": 0.14092140921409213,
-      "grad_norm": 0.30322559816019845,
+      "grad_norm": 2.377825424765151,
       "learning_rate": 2.9359792462981007e-05,
-      "loss": 0.599,
+      "loss": 0.5978,
       "step": 78
     },
     {
       "epoch": 0.14272809394760613,
-      "grad_norm": 0.29351194680487536,
+      "grad_norm": 0.4456515263682091,
       "learning_rate": 2.9333592086792113e-05,
-      "loss": 0.6075,
+      "loss": 0.6069,
       "step": 79
     },
     {
       "epoch": 0.14453477868112014,
-      "grad_norm": 0.2670107767188747,
+      "grad_norm": 0.34569754140915615,
       "learning_rate": 2.9306878453600382e-05,
-      "loss": 0.6161,
+      "loss": 0.6146,
       "step": 80
     },
     {
       "epoch": 0.14634146341463414,
-      "grad_norm": 0.2927366904918854,
+      "grad_norm": 0.41523607463693246,
       "learning_rate": 2.9279652519967105e-05,
-      "loss": 0.6004,
+      "loss": 0.5989,
       "step": 81
     },
     {
       "epoch": 0.14814814814814814,
-      "grad_norm": 0.28777141201555767,
+      "grad_norm": 0.4045393297763538,
       "learning_rate": 2.9251915260798024e-05,
-      "loss": 0.593,
+      "loss": 0.5913,
       "step": 82
     },
     {
       "epoch": 0.14995483288166214,
-      "grad_norm": 0.282262421377073,
+      "grad_norm": 0.4028600904179543,
       "learning_rate": 2.9223667669308395e-05,
-      "loss": 0.6038,
+      "loss": 0.6016,
       "step": 83
     },
     {
       "epoch": 0.15176151761517614,
-      "grad_norm": 0.24811154699761442,
+      "grad_norm": 0.34357038583025223,
       "learning_rate": 2.9194910756987464e-05,
-      "loss": 0.5976,
+      "loss": 0.5953,
       "step": 84
     },
     {
       "epoch": 0.15356820234869015,
-      "grad_norm": 0.27670190484589513,
+      "grad_norm": 0.37799582892290484,
       "learning_rate": 2.9165645553562215e-05,
-      "loss": 0.5997,
+      "loss": 0.5983,
       "step": 85
     },
     {
       "epoch": 0.15537488708220415,
-      "grad_norm": 0.24338487756299684,
+      "grad_norm": 1.4320099283126178,
       "learning_rate": 2.9135873106960525e-05,
-      "loss": 0.5994,
+      "loss": 0.5969,
       "step": 86
     },
     {
       "epoch": 0.15718157181571815,
-      "grad_norm": 0.3040191021283788,
+      "grad_norm": 0.35047306772001974,
       "learning_rate": 2.9105594483273603e-05,
-      "loss": 0.6012,
+      "loss": 0.6003,
       "step": 87
     },
     {
       "epoch": 0.15898825654923215,
-      "grad_norm": 0.26703924304748566,
+      "grad_norm": 0.4630513209073766,
       "learning_rate": 2.9074810766717865e-05,
-      "loss": 0.6058,
+      "loss": 0.6063,
       "step": 88
     },
     {
       "epoch": 0.16079494128274616,
-      "grad_norm": 0.2694312104706441,
+      "grad_norm": 0.38854101508245126,
       "learning_rate": 2.904352305959606e-05,
-      "loss": 0.602,
+      "loss": 0.6016,
       "step": 89
     },
     {
       "epoch": 0.16260162601626016,
-      "grad_norm": 0.29517928394037635,
+      "grad_norm": 0.40449057823694373,
       "learning_rate": 2.9011732482257835e-05,
-      "loss": 0.5944,
+      "loss": 0.5934,
       "step": 90
     },
     {
       "epoch": 0.16440831074977416,
-      "grad_norm": 0.25118801401605734,
+      "grad_norm": 0.43904838090142206,
       "learning_rate": 2.89794401730596e-05,
-      "loss": 0.5923,
+      "loss": 0.5921,
       "step": 91
     },
     {
       "epoch": 0.16621499548328816,
-      "grad_norm": 0.3348680034328015,
+      "grad_norm": 0.4396255315476391,
       "learning_rate": 2.894664728832377e-05,
-      "loss": 0.5958,
+      "loss": 0.5959,
       "step": 92
     },
     {
       "epoch": 0.16802168021680217,
-      "grad_norm": 0.25945936645525364,
+      "grad_norm": 0.38031370648914475,
       "learning_rate": 2.8913355002297367e-05,
-      "loss": 0.6025,
+      "loss": 0.6022,
       "step": 93
     },
     {
       "epoch": 0.16982836495031617,
-      "grad_norm": 0.31672412557275315,
+      "grad_norm": 0.523279763736358,
       "learning_rate": 2.887956450710995e-05,
-      "loss": 0.602,
+      "loss": 0.6028,
       "step": 94
     },
     {
       "epoch": 0.17163504968383017,
-      "grad_norm": 0.2631935513788724,
+      "grad_norm": 0.6704399767230238,
       "learning_rate": 2.8845277012730963e-05,
-      "loss": 0.5939,
+      "loss": 0.5942,
       "step": 95
     },
     {
       "epoch": 0.17344173441734417,
-      "grad_norm": 0.28309967073656406,
+      "grad_norm": 0.5066644836900498,
       "learning_rate": 2.8810493746926364e-05,
-      "loss": 0.5932,
+      "loss": 0.5933,
       "step": 96
     },
     {
       "epoch": 0.17524841915085818,
-      "grad_norm": 0.28939178950900774,
+      "grad_norm": 0.3646318333476554,
       "learning_rate": 2.87752159552147e-05,
-      "loss": 0.5993,
+      "loss": 0.5997,
       "step": 97
     },
     {
       "epoch": 0.17705510388437218,
-      "grad_norm": 0.3021445789644071,
+      "grad_norm": 0.376548770661227,
       "learning_rate": 2.87394449008225e-05,
-      "loss": 0.5924,
+      "loss": 0.5916,
       "step": 98
     },
     {
       "epoch": 0.17886178861788618,
-      "grad_norm": 0.2648161341085879,
+      "grad_norm": 0.4285200046228619,
       "learning_rate": 2.8703181864639013e-05,
-      "loss": 0.6058,
+      "loss": 0.6057,
       "step": 99
     },
     {
       "epoch": 0.18066847335140018,
-      "grad_norm": 0.28484263640616314,
+      "grad_norm": 0.33939531318747057,
       "learning_rate": 2.8666428145170385e-05,
-      "loss": 0.5789,
+      "loss": 0.5779,
       "step": 100
     },
     {
       "epoch": 0.18247515808491419,
-      "grad_norm": 0.2611846886724786,
+      "grad_norm": 0.3969014717910721,
       "learning_rate": 2.8629185058493116e-05,
-      "loss": 0.5949,
+      "loss": 0.5944,
       "step": 101
     },
     {
       "epoch": 0.1842818428184282,
-      "grad_norm": 0.33149916071305763,
+      "grad_norm": 0.36815480889738744,
       "learning_rate": 2.8591453938206985e-05,
-      "loss": 0.6134,
+      "loss": 0.6126,
       "step": 102
     },
     {
       "epoch": 0.1860885275519422,
-      "grad_norm": 0.28402400837841096,
+      "grad_norm": 0.3820598707301275,
       "learning_rate": 2.8553236135387247e-05,
-      "loss": 0.5914,
+      "loss": 0.5911,
       "step": 103
     },
     {
       "epoch": 0.1878952122854562,
-      "grad_norm": 0.28825320782588004,
+      "grad_norm": 0.4452739463450766,
       "learning_rate": 2.8514533018536286e-05,
-      "loss": 0.5897,
+      "loss": 0.5887,
       "step": 104
     },
     {
       "epoch": 0.1897018970189702,
-      "grad_norm": 0.2901395812522203,
+      "grad_norm": 0.37087706923300234,
       "learning_rate": 2.8475345973534605e-05,
-      "loss": 0.5822,
+      "loss": 0.5817,
       "step": 105
     },
     {
       "epoch": 0.1915085817524842,
-      "grad_norm": 0.25034051990006584,
+      "grad_norm": 0.6988372458559541,
       "learning_rate": 2.8435676403591193e-05,
-      "loss": 0.5846,
+      "loss": 0.584,
       "step": 106
     },
     {
       "epoch": 0.1933152664859982,
-      "grad_norm": 0.27868357694305945,
+      "grad_norm": 0.3800814296781821,
       "learning_rate": 2.8395525729193284e-05,
-      "loss": 0.5907,
+      "loss": 0.59,
       "step": 107
     },
     {
       "epoch": 0.1951219512195122,
-      "grad_norm": 0.257742952703987,
+      "grad_norm": 0.3993950018851527,
       "learning_rate": 2.835489538805548e-05,
-      "loss": 0.5888,
+      "loss": 0.589,
       "step": 108
     },
     {
       "epoch": 0.1969286359530262,
-      "grad_norm": 0.26397326166674545,
+      "grad_norm": 0.34149177022116556,
       "learning_rate": 2.8313786835068314e-05,
-      "loss": 0.5976,
+      "loss": 0.5973,
       "step": 109
     },
     {
       "epoch": 0.1987353206865402,
-      "grad_norm": 0.2721484510112116,
+      "grad_norm": 0.3199472139502249,
       "learning_rate": 2.8272201542246077e-05,
-      "loss": 0.5792,
+      "loss": 0.5791,
       "step": 110
     },
     {
       "epoch": 0.2005420054200542,
-      "grad_norm": 0.2374425057236497,
+      "grad_norm": 0.3608234766567561,
       "learning_rate": 2.8230140998674185e-05,
-      "loss": 0.599,
+      "loss": 0.5989,
       "step": 111
     },
     {
       "epoch": 0.2023486901535682,
-      "grad_norm": 0.24242218904343407,
+      "grad_norm": 0.3371671357477985,
       "learning_rate": 2.8187606710455807e-05,
-      "loss": 0.5921,
+      "loss": 0.5922,
       "step": 112
     },
     {
       "epoch": 0.2041553748870822,
-      "grad_norm": 0.2742052397132371,
+      "grad_norm": 0.36397935298555917,
       "learning_rate": 2.8144600200657953e-05,
-      "loss": 0.5965,
+      "loss": 0.5962,
       "step": 113
     },
     {
       "epoch": 0.20596205962059622,
-      "grad_norm": 0.2240931438132582,
+      "grad_norm": 0.44229892857297654,
       "learning_rate": 2.8101123009256946e-05,
-      "loss": 0.5792,
+      "loss": 0.5803,
       "step": 114
     },
     {
       "epoch": 0.20776874435411022,
-      "grad_norm": 0.2763456467888191,
+      "grad_norm": 0.3933110247729101,
       "learning_rate": 2.8057176693083253e-05,
-      "loss": 0.5958,
+      "loss": 0.5971,
       "step": 115
     },
     {
       "epoch": 0.20957542908762422,
-      "grad_norm": 0.2435434093260049,
+      "grad_norm": 0.35133090195312044,
       "learning_rate": 2.8012762825765763e-05,
-      "loss": 0.5891,
+      "loss": 0.59,
       "step": 116
     },
     {
       "epoch": 0.21138211382113822,
-      "grad_norm": 0.2769370185736096,
+      "grad_norm": 0.42414767478573207,
       "learning_rate": 2.7967882997675424e-05,
-      "loss": 0.6012,
+      "loss": 0.602,
       "step": 117
     },
     {
       "epoch": 0.21318879855465223,
-      "grad_norm": 0.2556991024086139,
+      "grad_norm": 0.3435364783148156,
       "learning_rate": 2.7922538815868287e-05,
-      "loss": 0.5879,
+      "loss": 0.5881,
       "step": 118
     },
     {
       "epoch": 0.21499548328816623,
-      "grad_norm": 0.2913165244476914,
+      "grad_norm": 0.33943051644040134,
       "learning_rate": 2.7876731904027994e-05,
-      "loss": 0.5784,
+      "loss": 0.5788,
       "step": 119
     },
     {
       "epoch": 0.21680216802168023,
-      "grad_norm": 0.29037318268167045,
+      "grad_norm": 0.3221999031967881,
       "learning_rate": 2.78304639024076e-05,
-      "loss": 0.5968,
+      "loss": 0.597,
       "step": 120
     },
     {
       "epoch": 0.2186088527551942,
-      "grad_norm": 0.24340049924612678,
+      "grad_norm": 0.7107243183485162,
       "learning_rate": 2.7783736467770863e-05,
       "loss": 0.5927,
       "step": 121
     },
     {
       "epoch": 0.2204155374887082,
-      "grad_norm": 0.29825692316659325,
+      "grad_norm": 0.4275366085321626,
       "learning_rate": 2.7736551273332908e-05,
-      "loss": 0.5972,
+      "loss": 0.5977,
       "step": 122
     },
     {
       "epoch": 0.2222222222222222,
-      "grad_norm": 0.2593832259683122,
+      "grad_norm": 0.34325044467207816,
       "learning_rate": 2.7688910008700305e-05,
-      "loss": 0.5876,
+      "loss": 0.5883,
       "step": 123
     },
     {
       "epoch": 0.2240289069557362,
-      "grad_norm": 0.27547771011759425,
+      "grad_norm": 0.3503193328962726,
       "learning_rate": 2.764081437981059e-05,
-      "loss": 0.5742,
+      "loss": 0.5747,
       "step": 124
     },
     {
       "epoch": 0.22583559168925021,
-      "grad_norm": 0.27158658907185634,
+      "grad_norm": 0.39807179938908716,
       "learning_rate": 2.7592266108871158e-05,
-      "loss": 0.5812,
+      "loss": 0.5816,
       "step": 125
     },
     {
       "epoch": 0.22764227642276422,
-      "grad_norm": 0.2715497634460973,
+      "grad_norm": 0.36593993813666903,
       "learning_rate": 2.754326693429761e-05,
-      "loss": 0.5931,
+      "loss": 0.5936,
       "step": 126
     },
     {
       "epoch": 0.22944896115627822,
-      "grad_norm": 0.24142985625388147,
+      "grad_norm": 0.3871443516927523,
       "learning_rate": 2.7493818610651493e-05,
-      "loss": 0.5859,
+      "loss": 0.5875,
       "step": 127
     },
     {
       "epoch": 0.23125564588979222,
-      "grad_norm": 0.2530474955317245,
+      "grad_norm": 0.3606586360294697,
       "learning_rate": 2.744392290857747e-05,
-      "loss": 0.5857,
+      "loss": 0.5866,
       "step": 128
     },
     {
       "epoch": 0.23306233062330622,
-      "grad_norm": 0.2567731311347638,
+      "grad_norm": 0.38342216198991025,
       "learning_rate": 2.7393581614739924e-05,
-      "loss": 0.5803,
+      "loss": 0.5813,
       "step": 129
     },
     {
       "epoch": 0.23486901535682023,
-      "grad_norm": 0.2628073110459229,
+      "grad_norm": 0.37465039657790306,
       "learning_rate": 2.7342796531758984e-05,
-      "loss": 0.5987,
+      "loss": 0.6004,
       "step": 130
     },
     {
       "epoch": 0.23667570009033423,
-      "grad_norm": 0.23510376935644756,
+      "grad_norm": 0.3764086380678446,
       "learning_rate": 2.729156947814598e-05,
-      "loss": 0.5926,
+      "loss": 0.5943,
       "step": 131
     },
     {
       "epoch": 0.23848238482384823,
-      "grad_norm": 0.23576954058784624,
+      "grad_norm": 0.7637683242409493,
       "learning_rate": 2.7239902288238297e-05,
-      "loss": 0.5721,
+      "loss": 0.5737,
       "step": 132
     },
     {
       "epoch": 0.24028906955736223,
-      "grad_norm": 0.24964231366212492,
+      "grad_norm": 0.3560589890890579,
       "learning_rate": 2.7187796812133733e-05,
-      "loss": 0.5873,
+      "loss": 0.5887,
       "step": 133
     },
     {
       "epoch": 0.24209575429087624,
-      "grad_norm": 0.23911643013915115,
+      "grad_norm": 0.3825941152602638,
       "learning_rate": 2.7135254915624213e-05,
-      "loss": 0.5878,
+      "loss": 0.59,
       "step": 134
     },
     {
       "epoch": 0.24390243902439024,
-      "grad_norm": 0.23601189225311667,
+      "grad_norm": 0.36389498688783184,
       "learning_rate": 2.708227848012901e-05,
-      "loss": 0.585,
+      "loss": 0.5876,
       "step": 135
     },
     {
       "epoch": 0.24570912375790424,
-      "grad_norm": 0.28061149806949875,
+      "grad_norm": 0.37878840510853407,
       "learning_rate": 2.7028869402627357e-05,
-      "loss": 0.5889,
+      "loss": 0.591,
       "step": 136
     },
     {
       "epoch": 0.24751580849141824,
-      "grad_norm": 0.2537605160128264,
+      "grad_norm": 0.3899552482148334,
       "learning_rate": 2.6975029595590523e-05,
-      "loss": 0.5726,
+      "loss": 0.5742,
       "step": 137
     },
     {
       "epoch": 0.24932249322493225,
-      "grad_norm": 0.2631279788827005,
+      "grad_norm": 0.3807471778199868,
       "learning_rate": 2.6920760986913332e-05,
-      "loss": 0.5901,
+      "loss": 0.592,
       "step": 138
     },
     {
       "epoch": 0.25112917795844625,
-      "grad_norm": 0.24684546421512338,
+      "grad_norm": 0.41701742101917744,
       "learning_rate": 2.6866065519845124e-05,
-      "loss": 0.5932,
+      "loss": 0.5953,
       "step": 139
     },
     {
       "epoch": 0.2529358626919603,
-      "grad_norm": 0.25108587954115535,
+      "grad_norm": 0.39189298513945836,
       "learning_rate": 2.681094515292018e-05,
-      "loss": 0.5927,
+      "loss": 0.5943,
       "step": 140
     },
     {
       "epoch": 0.25474254742547425,
-      "grad_norm": 0.2759732073497199,
+      "grad_norm": 0.5240993515099046,
       "learning_rate": 2.6755401859887598e-05,
-      "loss": 0.573,
+      "loss": 0.5755,
       "step": 141
     },
     {
       "epoch": 0.2565492321589883,
-      "grad_norm": 0.2927986426856055,
+      "grad_norm": 0.37512137985409255,
       "learning_rate": 2.6699437629640595e-05,
-      "loss": 0.5934,
+      "loss": 0.5956,
       "step": 142
     },
     {
       "epoch": 0.25835591689250226,
-      "grad_norm": 0.2915811926744858,
+      "grad_norm": 0.3707219843953511,
       "learning_rate": 2.6643054466145297e-05,
-      "loss": 0.5846,
+      "loss": 0.5869,
       "step": 143
     },
     {
       "epoch": 0.2601626016260163,
-      "grad_norm": 0.25639674407942853,
+      "grad_norm": 0.3555027971231102,
       "learning_rate": 2.6586254388368995e-05,
-      "loss": 0.5733,
+      "loss": 0.5752,
       "step": 144
     },
     {
       "epoch": 0.26196928635953026,
-      "grad_norm": 0.25466281147535325,
+      "grad_norm": 0.3026146177016955,
       "learning_rate": 2.652903943020783e-05,
-      "loss": 0.5862,
+      "loss": 0.5883,
       "step": 145
     },
     {
       "epoch": 0.26377597109304424,
-      "grad_norm": 0.24151461009286013,
+      "grad_norm": 0.41020681896565947,
       "learning_rate": 2.647141164041398e-05,
-      "loss": 0.5904,
+      "loss": 0.5937,
       "step": 146
     },
     {
       "epoch": 0.26558265582655827,
-      "grad_norm": 0.259822033974557,
+      "grad_norm": 0.36661807722028844,
       "learning_rate": 2.641337308252228e-05,
-      "loss": 0.581,
+      "loss": 0.583,
       "step": 147
     },
     {
       "epoch": 0.26738934056007224,
-      "grad_norm": 0.24532475377494284,
+      "grad_norm": 0.32357318798570317,
       "learning_rate": 2.6354925834776346e-05,
-      "loss": 0.5817,
+      "loss": 0.5841,
       "step": 148
     },
     {
       "epoch": 0.26919602529358627,
-      "grad_norm": 0.2555788655192275,
+      "grad_norm": 0.4240562481643986,
       "learning_rate": 2.6296071990054167e-05,
-      "loss": 0.579,
+      "loss": 0.5812,
       "step": 149
     },
     {
       "epoch": 0.27100271002710025,
-      "grad_norm": 0.2534339164884998,
+      "grad_norm": 0.35429923851701156,
       "learning_rate": 2.6236813655793123e-05,
-      "loss": 0.5886,
+      "loss": 0.5905,
       "step": 150
     },
     {
       "epoch": 0.2728093947606143,
-      "grad_norm": 0.25366511050468843,
+      "grad_norm": 0.3981067430459517,
       "learning_rate": 2.617715295391457e-05,
-      "loss": 0.5707,
+      "loss": 0.572,
       "step": 151
     },
     {
       "epoch": 0.27461607949412825,
-      "grad_norm": 0.26118391438335414,
+      "grad_norm": 0.516641806924778,
       "learning_rate": 2.6117092020747824e-05,
-      "loss": 0.5779,
+      "loss": 0.5802,
       "step": 152
     },
     {
       "epoch": 0.2764227642276423,
-      "grad_norm": 0.31049055796662123,
+      "grad_norm": 0.3349831187986836,
       "learning_rate": 2.6056633006953677e-05,
-      "loss": 0.5783,
+      "loss": 0.581,
       "step": 153
     },
     {
       "epoch": 0.27822944896115626,
-      "grad_norm": 0.2611788471953917,
+      "grad_norm": 0.32079723720771636,
       "learning_rate": 2.5995778077447393e-05,
-      "loss": 0.5823,
+      "loss": 0.5847,
       "step": 154
     },
     {
       "epoch": 0.2800361336946703,
-      "grad_norm": 0.3259190805002546,
+      "grad_norm": 0.4491293903796211,
       "learning_rate": 2.5934529411321174e-05,
-      "loss": 0.5743,
+      "loss": 0.577,
       "step": 155
     },
     {
       "epoch": 0.28184281842818426,
-      "grad_norm": 0.2222204544644696,
+      "grad_norm": 0.3338078899554527,
       "learning_rate": 2.587288920176613e-05,
-      "loss": 0.575,
+      "loss": 0.577,
       "step": 156
     },
     {
       "epoch": 0.2836495031616983,
-      "grad_norm": 0.29042132848143276,
+      "grad_norm": 0.337095193899371,
       "learning_rate": 2.581085965599375e-05,
-      "loss": 0.5877,
+      "loss": 0.5905,
       "step": 157
     },
     {
       "epoch": 0.28545618789521227,
-      "grad_norm": 0.2577824036794116,
+      "grad_norm": 0.3241526231317586,
       "learning_rate": 2.5748442995156882e-05,
-      "loss": 0.5854,
+      "loss": 0.5872,
       "step": 158
     },
     {
       "epoch": 0.2872628726287263,
-      "grad_norm": 0.24041412490024394,
+      "grad_norm": 0.32809237233012317,
       "learning_rate": 2.5685641454270172e-05,
-      "loss": 0.566,
+      "loss": 0.5683,
       "step": 159
     },
     {
       "epoch": 0.28906955736224027,
-      "grad_norm": 0.33137613211077116,
+      "grad_norm": 0.33268577227395,
       "learning_rate": 2.5622457282130046e-05,
-      "loss": 0.5753,
+      "loss": 0.5778,
       "step": 160
     },
     {
       "epoch": 0.2908762420957543,
-      "grad_norm": 0.2252657452973701,
+      "grad_norm": 0.4132273947780692,
       "learning_rate": 2.5558892741234173e-05,
-      "loss": 0.5827,
+      "loss": 0.5861,
       "step": 161
     },
     {
       "epoch": 0.2926829268292683,
-      "grad_norm": 0.353225035311219,
+      "grad_norm": 0.33902072712237397,
       "learning_rate": 2.5494950107700482e-05,
-      "loss": 0.5692,
+      "loss": 0.5709,
       "step": 162
     },
     {
       "epoch": 0.2944896115627823,
-      "grad_norm": 0.2549710140850951,
+      "grad_norm": 0.3587450885360311,
       "learning_rate": 2.5430631671185616e-05,
-      "loss": 0.5775,
+      "loss": 0.5791,
       "step": 163
     },
     {
       "epoch": 0.2962962962962963,
-      "grad_norm": 0.320108207545845,
+      "grad_norm": 0.32948214359444944,
       "learning_rate": 2.5365939734802973e-05,
-      "loss": 0.5871,
+      "loss": 0.5892,
       "step": 164
     },
     {
       "epoch": 0.2981029810298103,
-      "grad_norm": 0.27118178617195304,
+      "grad_norm": 0.36655920600202585,
       "learning_rate": 2.5300876615040223e-05,
-      "loss": 0.5777,
+      "loss": 0.5802,
       "step": 165
     },
     {
       "epoch": 0.2999096657633243,
-      "grad_norm": 0.2732065296628375,
+      "grad_norm": 0.36077426576274396,
       "learning_rate": 2.523544464167637e-05,
-      "loss": 0.5842,
+      "loss": 0.5868,
       "step": 166
     },
     {
       "epoch": 0.3017163504968383,
-      "grad_norm": 0.2440418264084237,
+      "grad_norm": 0.35802150689356144,
       "learning_rate": 2.5169646157698313e-05,
-      "loss": 0.583,
+      "loss": 0.5857,
       "step": 167
     },
     {
       "epoch": 0.3035230352303523,
-      "grad_norm": 0.2400624724182379,
+      "grad_norm": 0.3262888729605983,
       "learning_rate": 2.5103483519216964e-05,
-      "loss": 0.5813,
+      "loss": 0.5838,
       "step": 168
     },
     {
       "epoch": 0.3053297199638663,
-      "grad_norm": 0.25905283588713374,
+      "grad_norm": 0.3362096574051899,
       "learning_rate": 2.5036959095382875e-05,
-      "loss": 0.574,
+      "loss": 0.5755,
       "step": 169
     },
     {
       "epoch": 0.3071364046973803,
-      "grad_norm": 0.2591290671871167,
+      "grad_norm": 0.4060212080395927,
       "learning_rate": 2.4970075268301388e-05,
-      "loss": 0.5877,
+      "loss": 0.5893,
       "step": 170
     },
     {
       "epoch": 0.3089430894308943,
-      "grad_norm": 0.2235510897395504,
+      "grad_norm": 0.2886329167833011,
       "learning_rate": 2.4902834432947353e-05,
-      "loss": 0.5706,
+      "loss": 0.5733,
       "step": 171
     },
     {
       "epoch": 0.3107497741644083,
-      "grad_norm": 0.23760061910249497,
+      "grad_norm": 0.3808717667450013,
       "learning_rate": 2.4835238997079382e-05,
-      "loss": 0.5771,
+      "loss": 0.5795,
       "step": 172
     },
     {
       "epoch": 0.31255645889792233,
-      "grad_norm": 0.2617951363502666,
+      "grad_norm": 0.3072237334054852,
       "learning_rate": 2.4767291381153603e-05,
-      "loss": 0.5807,
+      "loss": 0.5832,
       "step": 173
     },
     {
       "epoch": 0.3143631436314363,
-      "grad_norm": 0.23091611449448254,
+      "grad_norm": 0.2965837108398263,
       "learning_rate": 2.4698994018236994e-05,
-      "loss": 0.5721,
+      "loss": 0.5747,
       "step": 174
     },
     {
       "epoch": 0.31616982836495033,
-      "grad_norm": 0.25262572503013164,
+      "grad_norm": 0.3384696066427267,
       "learning_rate": 2.4630349353920284e-05,
-      "loss": 0.5767,
+      "loss": 0.5784,
       "step": 175
     },
     {
       "epoch": 0.3179765130984643,
-      "grad_norm": 0.23055115622552902,
+      "grad_norm": 0.30615079028521514,
       "learning_rate": 2.4561359846230346e-05,
-      "loss": 0.5906,
+      "loss": 0.5923,
       "step": 176
     },
     {
       "epoch": 0.31978319783197834,
-      "grad_norm": 0.21968694104114944,
+      "grad_norm": 0.2888166567825809,
       "learning_rate": 2.4492027965542217e-05,
-      "loss": 0.5657,
+      "loss": 0.5679,
       "step": 177
     },
     {
       "epoch": 0.3215898825654923,
-      "grad_norm": 0.2525014182918365,
+      "grad_norm": 0.30385385604927617,
       "learning_rate": 2.44223561944906e-05,
-      "loss": 0.5724,
+      "loss": 0.5744,
       "step": 178
     },
     {
       "epoch": 0.32339656729900634,
-      "grad_norm": 0.2407328065060556,
+      "grad_norm": 0.28021720457684524,
       "learning_rate": 2.4352347027881003e-05,
-      "loss": 0.5825,
+      "loss": 0.5846,
       "step": 179
     },
     {
       "epoch": 0.3252032520325203,
-      "grad_norm": 0.2547387307288275,
+      "grad_norm": 0.30261488499331,
       "learning_rate": 2.4282002972600382e-05,
-      "loss": 0.5908,
+      "loss": 0.5932,
       "step": 180
     },
     {
       "epoch": 0.32700993676603435,
-      "grad_norm": 0.2263495601766316,
+      "grad_norm": 0.2851073248221503,
       "learning_rate": 2.4211326547527377e-05,
-      "loss": 0.5688,
+      "loss": 0.5706,
       "step": 181
     },
     {
       "epoch": 0.3288166214995483,
-      "grad_norm": 0.2789402543570071,
+      "grad_norm": 0.3111115931546976,
       "learning_rate": 2.4140320283442125e-05,
-      "loss": 0.579,
+      "loss": 0.5813,
       "step": 182
     },
     {
       "epoch": 0.33062330623306235,
-      "grad_norm": 0.2567309540458022,
+      "grad_norm": 0.3143186202532528,
       "learning_rate": 2.4068986722935625e-05,
-      "loss": 0.5893,
+      "loss": 0.5924,
       "step": 183
     },
     {
       "epoch": 0.3324299909665763,
-      "grad_norm": 0.24182513399667138,
+      "grad_norm": 0.35464611022112247,
       "learning_rate": 2.3997328420318705e-05,
-      "loss": 0.5858,
+      "loss": 0.5885,
       "step": 184
     },
     {
       "epoch": 0.33423667570009036,
-      "grad_norm": 0.2675738732389291,
+      "grad_norm": 0.3119101078690158,
       "learning_rate": 2.3925347941530556e-05,
-      "loss": 0.5822,
+      "loss": 0.585,
       "step": 185
     },
     {
       "epoch": 0.33604336043360433,
-      "grad_norm": 0.23428952659343405,
+      "grad_norm": 0.3362124015760118,
       "learning_rate": 2.3853047864046843e-05,
-      "loss": 0.5779,
+      "loss": 0.5805,
       "step": 186
     },
     {
       "epoch": 0.33785004516711836,
-      "grad_norm": 0.26339154899807604,
+      "grad_norm": 0.31155910265392833,
       "learning_rate": 2.3780430776787413e-05,
-      "loss": 0.5747,
+      "loss": 0.5775,
       "step": 187
     },
     {
       "epoch": 0.33965672990063234,
-      "grad_norm": 0.21349849308750005,
+      "grad_norm": 0.30082608710665315,
       "learning_rate": 2.3707499280023604e-05,
-      "loss": 0.5747,
+      "loss": 0.5764,
       "step": 188
     },
     {
       "epoch": 0.34146341463414637,
-      "grad_norm": 0.26554041110868815,
+      "grad_norm": 0.3051156622197956,
       "learning_rate": 2.3634255985285104e-05,
-      "loss": 0.5825,
+      "loss": 0.5845,
       "step": 189
     },
     {
       "epoch": 0.34327009936766034,
-      "grad_norm": 0.23538164196617103,
+      "grad_norm": 0.37884369981203814,
       "learning_rate": 2.356070351526648e-05,
-      "loss": 0.5708,
+      "loss": 0.5736,
       "step": 190
     },
     {
       "epoch": 0.34507678410117437,
-      "grad_norm": 0.24516795740455555,
+      "grad_norm": 0.32478176061158476,
       "learning_rate": 2.348684450373322e-05,
-      "loss": 0.5691,
+      "loss": 0.5721,
       "step": 191
     },
     {
       "epoch": 0.34688346883468835,
-      "grad_norm": 0.2320939063694949,
+      "grad_norm": 0.8776768851811568,
       "learning_rate": 2.3412681595427467e-05,
-      "loss": 0.5761,
+      "loss": 0.579,
       "step": 192
     },
     {
       "epoch": 0.3486901535682023,
-      "grad_norm": 0.24367157013907606,
+      "grad_norm": 0.4610655284279262,
       "learning_rate": 2.3338217445973268e-05,
-      "loss": 0.5776,
+      "loss": 0.5806,
       "step": 193
     },
     {
       "epoch": 0.35049683830171635,
-      "grad_norm": 0.2383952654545809,
+      "grad_norm": 0.32423892822372186,
       "learning_rate": 2.3263454721781537e-05,
-      "loss": 0.5793,
+      "loss": 0.5824,
       "step": 194
     },
     {
       "epoch": 0.3523035230352303,
-      "grad_norm": 0.21970860044862534,
+      "grad_norm": 0.3305614412713182,
       "learning_rate": 2.318839609995453e-05,
-      "loss": 0.5625,
+      "loss": 0.5655,
       "step": 195
     },
     {
       "epoch": 0.35411020776874436,
-      "grad_norm": 0.22393266445368756,
+      "grad_norm": 0.3088383542670599,
       "learning_rate": 2.3113044268189995e-05,
-      "loss": 0.5758,
+      "loss": 0.579,
       "step": 196
     },
     {
       "epoch": 0.35591689250225833,
-      "grad_norm": 0.20965253283535495,
+      "grad_norm": 0.3087829310275476,
       "learning_rate": 2.303740192468495e-05,
-      "loss": 0.5754,
+      "loss": 0.5782,
       "step": 197
     },
     {
       "epoch": 0.35772357723577236,
-      "grad_norm": 0.2134213495454822,
+      "grad_norm": 0.5804440405125408,
       "learning_rate": 2.2961471778039045e-05,
-      "loss": 0.5521,
+      "loss": 0.5552,
       "step": 198
     },
     {
       "epoch": 0.35953026196928634,
-      "grad_norm": 0.2077804399624621,
+      "grad_norm": 0.3215405528760461,
       "learning_rate": 2.288525654715757e-05,
-      "loss": 0.5758,
+      "loss": 0.5797,
       "step": 199
     },
     {
       "epoch": 0.36133694670280037,
-      "grad_norm": 0.21089298206975676,
+      "grad_norm": 0.34958938754043156,
       "learning_rate": 2.280875896115413e-05,
-      "loss": 0.5645,
+      "loss": 0.5676,
       "step": 200
     },
     {
       "epoch": 0.36314363143631434,
-      "grad_norm": 0.22751400705060473,
+      "grad_norm": 0.33248083840336357,
       "learning_rate": 2.2731981759252876e-05,
-      "loss": 0.5595,
+      "loss": 0.5624,
       "step": 201
     },
     {
       "epoch": 0.36495031616982837,
-      "grad_norm": 0.22071355320317226,
+      "grad_norm": 0.26750582916819504,
       "learning_rate": 2.2654927690690445e-05,
-      "loss": 0.5858,
+      "loss": 0.5887,
       "step": 202
     },
     {
       "epoch": 0.36675700090334235,
-      "grad_norm": 0.2269131246517602,
+      "grad_norm": 0.3306138445776805,
       "learning_rate": 2.257759951461752e-05,
-      "loss": 0.5733,
+      "loss": 0.5765,
       "step": 203
     },
     {
       "epoch": 0.3685636856368564,
-      "grad_norm": 0.27195813360203375,
+      "grad_norm": 0.3309054658853368,
       "learning_rate": 2.25e-05,
-      "loss": 0.5855,
+      "loss": 0.588,
       "step": 204
     },
     {
       "epoch": 0.37037037037037035,
-      "grad_norm": 0.24074584104894936,
+      "grad_norm": 0.3069934073052549,
       "learning_rate": 2.24221319255199e-05,
-      "loss": 0.5761,
+      "loss": 0.5801,
       "step": 205
     },
     {
       "epoch": 0.3721770551038844,
-      "grad_norm": 0.2643169067368919,
+      "grad_norm": 0.36963521242301206,
       "learning_rate": 2.234399807947579e-05,
-      "loss": 0.5632,
+      "loss": 0.5669,
       "step": 206
     },
     {
       "epoch": 0.37398373983739835,
-      "grad_norm": 0.2602452192335381,
+      "grad_norm": 0.2842636472789742,
       "learning_rate": 2.2265601259683e-05,
-      "loss": 0.5713,
+      "loss": 0.5746,
       "step": 207
     },
     {
       "epoch": 0.3757904245709124,
-      "grad_norm": 0.22491768260975722,
+      "grad_norm": 0.2967300023743323,
       "learning_rate": 2.2186944273373426e-05,
-      "loss": 0.5589,
+      "loss": 0.5621,
       "step": 208
     },
     {
       "epoch": 0.37759710930442636,
-      "grad_norm": 0.23899265916988002,
+      "grad_norm": 0.30164540888223945,
       "learning_rate": 2.210802993709498e-05,
-      "loss": 0.5799,
+      "loss": 0.5841,
       "step": 209
     },
     {
       "epoch": 0.3794037940379404,
-      "grad_norm": 0.28185732149789655,
+      "grad_norm": 0.3020950057384688,
       "learning_rate": 2.202886107661078e-05,
-      "loss": 0.576,
+      "loss": 0.5803,
       "step": 210
     },
     {
       "epoch": 0.38121047877145436,
-      "grad_norm": 0.20671343969870537,
+      "grad_norm": 0.30358643623890824,
       "learning_rate": 2.1949440526797928e-05,
-      "loss": 0.5761,
+      "loss": 0.5798,
       "step": 211
     },
     {
       "epoch": 0.3830171635049684,
-      "grad_norm": 0.245945113158813,
+      "grad_norm": 0.35616750454359836,
       "learning_rate": 2.1869771131546015e-05,
-      "loss": 0.5834,
+      "loss": 0.587,
       "step": 212
     },
     {
       "epoch": 0.38482384823848237,
-      "grad_norm": 0.2051130167210825,
+      "grad_norm": 0.27878546685216377,
       "learning_rate": 2.178985574365529e-05,
-      "loss": 0.5837,
+      "loss": 0.5879,
       "step": 213
     },
     {
       "epoch": 0.3866305329719964,
-      "grad_norm": 0.25952305018947347,
+      "grad_norm": 0.32699494080692665,
       "learning_rate": 2.170969722473449e-05,
-      "loss": 0.57,
+      "loss": 0.5738,
       "step": 214
     },
     {
       "epoch": 0.3884372177055104,
-      "grad_norm": 0.24249808093673167,
+      "grad_norm": 0.31103539696417715,
       "learning_rate": 2.1629298445098403e-05,
-      "loss": 0.5966,
+      "loss": 0.6008,
       "step": 215
     },
     {
       "epoch": 0.3902439024390244,
-      "grad_norm": 0.237872387820411,
+      "grad_norm": 0.32794619585958185,
       "learning_rate": 2.154866228366505e-05,
-      "loss": 0.5856,
+      "loss": 0.5886,
       "step": 216
     },
     {
       "epoch": 0.3920505871725384,
-      "grad_norm": 0.20572539636470738,
+      "grad_norm": 0.36007383306015284,
       "learning_rate": 2.146779162785263e-05,
-      "loss": 0.5857,
+      "loss": 0.5896,
       "step": 217
     },
     {
       "epoch": 0.3938572719060524,
-      "grad_norm": 0.22849509908427684,
+      "grad_norm": 0.3059835677015186,
       "learning_rate": 2.138668937347609e-05,
-      "loss": 0.5775,
+      "loss": 0.5808,
       "step": 218
     },
     {
       "epoch": 0.3956639566395664,
-      "grad_norm": 0.21171703597428448,
+      "grad_norm": 0.33519811127594074,
       "learning_rate": 2.1305358424643484e-05,
-      "loss": 0.572,
+      "loss": 0.5758,
       "step": 219
     },
     {
       "epoch": 0.3974706413730804,
-      "grad_norm": 0.2333174990121186,
+      "grad_norm": 0.3173280187406549,
       "learning_rate": 2.1223801693651927e-05,
-      "loss": 0.5746,
+      "loss": 0.5778,
       "step": 220
     },
     {
       "epoch": 0.3992773261065944,
-      "grad_norm": 0.21728105831361583,
+      "grad_norm": 0.3249654059010307,
       "learning_rate": 2.114202210088336e-05,
-      "loss": 0.5633,
+      "loss": 0.5665,
       "step": 221
     },
     {
       "epoch": 0.4010840108401084,
-      "grad_norm": 0.2232469073657137,
+      "grad_norm": 0.315983288828985,
       "learning_rate": 2.106002257469993e-05,
-      "loss": 0.5805,
+      "loss": 0.5841,
       "step": 222
     },
     {
       "epoch": 0.4028906955736224,
-      "grad_norm": 0.23039919622981786,
+      "grad_norm": 0.317439517236036,
       "learning_rate": 2.0977806051339172e-05,
-      "loss": 0.5631,
+      "loss": 0.5661,
       "step": 223
     },
     {
       "epoch": 0.4046973803071364,
-      "grad_norm": 0.24767592985180426,
+      "grad_norm": 0.31483264135931766,
       "learning_rate": 2.0895375474808857e-05,
-      "loss": 0.5681,
+      "loss": 0.5714,
       "step": 224
     },
     {
       "epoch": 0.4065040650406504,
-      "grad_norm": 0.22933111893390376,
+      "grad_norm": 0.2948623418665935,
       "learning_rate": 2.0812733796781544e-05,
-      "loss": 0.5728,
+      "loss": 0.5765,
       "step": 225
     },
     {
       "epoch": 0.4083107497741644,
-      "grad_norm": 0.24263450877516124,
+      "grad_norm": 0.2966240742398221,
       "learning_rate": 2.0729883976488936e-05,
-      "loss": 0.5792,
+      "loss": 0.5831,
       "step": 226
     },
     {
       "epoch": 0.4101174345076784,
-      "grad_norm": 0.24644389391459087,
+      "grad_norm": 0.3938546298314616,
       "learning_rate": 2.064682898061588e-05,
-      "loss": 0.5836,
+      "loss": 0.5872,
       "step": 227
     },
     {
       "epoch": 0.41192411924119243,
-      "grad_norm": 0.23999077805752708,
+      "grad_norm": 0.2823560010807599,
       "learning_rate": 2.0563571783194146e-05,
-      "loss": 0.5697,
+      "loss": 0.573,
       "step": 228
     },
     {
       "epoch": 0.4137308039747064,
-      "grad_norm": 0.20338232613186608,
+      "grad_norm": 0.3226289122791917,
       "learning_rate": 2.0480115365495928e-05,
-      "loss": 0.5723,
+      "loss": 0.5759,
       "step": 229
     },
     {
       "epoch": 0.41553748870822044,
-      "grad_norm": 0.21451640836421845,
+      "grad_norm": 0.3116964943536504,
       "learning_rate": 2.0396462715927107e-05,
-      "loss": 0.5748,
+      "loss": 0.5782,
       "step": 230
     },
     {
       "epoch": 0.4173441734417344,
-      "grad_norm": 0.23674092702209837,
+      "grad_norm": 0.33810521611187627,
       "learning_rate": 2.0312616829920222e-05,
-      "loss": 0.5737,
+      "loss": 0.5771,
       "step": 231
     },
     {
       "epoch": 0.41915085817524844,
-      "grad_norm": 0.2040271869891673,
+      "grad_norm": 0.3046098546864739,
       "learning_rate": 2.022858070982723e-05,
-      "loss": 0.5768,
+      "loss": 0.5802,
       "step": 232
     },
     {
       "epoch": 0.4209575429087624,
-      "grad_norm": 0.23013110791835603,
+      "grad_norm": 0.2740022073595928,
       "learning_rate": 2.0144357364811973e-05,
-      "loss": 0.5586,
+      "loss": 0.5615,
       "step": 233
     },
     {
       "epoch": 0.42276422764227645,
-      "grad_norm": 0.23251714747465213,
+      "grad_norm": 0.2913067130533732,
       "learning_rate": 2.0059949810742452e-05,
-      "loss": 0.5744,
+      "loss": 0.5771,
       "step": 234
     },
     {
       "epoch": 0.4245709123757904,
-      "grad_norm": 0.24113338282908972,
+      "grad_norm": 0.2762685594049124,
       "learning_rate": 1.997536107008281e-05,
-      "loss": 0.5696,
+      "loss": 0.5725,
       "step": 235
     },
     {
       "epoch": 0.42637759710930445,
-      "grad_norm": 0.217275693069977,
+      "grad_norm": 0.5121383274688709,
       "learning_rate": 1.9890594171785128e-05,
-      "loss": 0.5807,
+      "loss": 0.5843,
       "step": 236
     },
     {
       "epoch": 0.4281842818428184,
-      "grad_norm": 0.20589216064741256,
+      "grad_norm": 0.3344358767381985,
       "learning_rate": 1.9805652151180945e-05,
-      "loss": 0.5804,
+      "loss": 0.5848,
       "step": 237
     },
     {
       "epoch": 0.42999096657633246,
-      "grad_norm": 0.22359306351640823,
+      "grad_norm": 0.31383841038930804,
       "learning_rate": 1.972053804987258e-05,
-      "loss": 0.5889,
+      "loss": 0.5924,
       "step": 238
     },
     {
       "epoch": 0.43179765130984643,
-      "grad_norm": 0.20375725610745427,
+      "grad_norm": 0.32911204189458476,
       "learning_rate": 1.963525491562421e-05,
-      "loss": 0.5701,
+      "loss": 0.573,
       "step": 239
     },
     {
       "epoch": 0.43360433604336046,
-      "grad_norm": 0.20454609097134233,
+      "grad_norm": 0.3543469488051476,
       "learning_rate": 1.954980580225275e-05,
-      "loss": 0.5702,
+      "loss": 0.5741,
       "step": 240
     },
     {
       "epoch": 0.43541102077687444,
-      "grad_norm": 0.19497163244699564,
+      "grad_norm": 0.2769864406337721,
       "learning_rate": 1.946419376951848e-05,
-      "loss": 0.5693,
+      "loss": 0.5736,
       "step": 241
     },
     {
       "epoch": 0.4372177055103884,
-      "grad_norm": 0.21551459772294188,
+      "grad_norm": 0.2772504692036322,
       "learning_rate": 1.9378421883015505e-05,
-      "loss": 0.57,
+      "loss": 0.5737,
       "step": 242
     },
     {
       "epoch": 0.43902439024390244,
-      "grad_norm": 0.23200634978654613,
+      "grad_norm": 0.27865104121234274,
       "learning_rate": 1.9292493214061953e-05,
-      "loss": 0.5741,
+      "loss": 0.5787,
       "step": 243
     },
     {
       "epoch": 0.4408310749774164,
-      "grad_norm": 0.1884963011054949,
+      "grad_norm": 0.33872737991020524,
       "learning_rate": 1.9206410839590042e-05,
-      "loss": 0.5804,
+      "loss": 0.5849,
       "step": 244
     },
     {
       "epoch": 0.44263775971093045,
-      "grad_norm": 0.25404075829436523,
+      "grad_norm": 0.29138500205569745,
       "learning_rate": 1.9120177842035853e-05,
-      "loss": 0.5726,
+      "loss": 0.5769,
       "step": 245
     },
     {
       "epoch": 0.4444444444444444,
-      "grad_norm": 0.22097319274710464,
+      "grad_norm": 0.33436184779749656,
       "learning_rate": 1.9033797309228984e-05,
-      "loss": 0.5794,
+      "loss": 0.584,
       "step": 246
     },
     {
       "epoch": 0.44625112917795845,
-      "grad_norm": 0.19246576672187318,
+      "grad_norm": 0.3402673798099814,
       "learning_rate": 1.8947272334281977e-05,
-      "loss": 0.5616,
+      "loss": 0.5654,
       "step": 247
     },
     {
       "epoch": 0.4480578139114724,
-      "grad_norm": 0.22700593999033858,
+      "grad_norm": 0.30708428061855253,
       "learning_rate": 1.8860606015479537e-05,
-      "loss": 0.5801,
+      "loss": 0.5846,
       "step": 248
     },
     {
       "epoch": 0.44986449864498645,
-      "grad_norm": 0.1934681513313933,
+      "grad_norm": 0.31233742569187306,
       "learning_rate": 1.877380145616763e-05,
-      "loss": 0.5632,
+      "loss": 0.5679,
       "step": 249
     },
     {
       "epoch": 0.45167118337850043,
-      "grad_norm": 0.21167453440737605,
+      "grad_norm": 0.3506675404562718,
       "learning_rate": 1.868686176464232e-05,
-      "loss": 0.5781,
+      "loss": 0.5836,
       "step": 250
     },
     {
       "epoch": 0.45347786811201446,
-      "grad_norm": 0.23469761452053498,
+      "grad_norm": 0.29769387160985,
       "learning_rate": 1.8599790054038487e-05,
-      "loss": 0.5653,
+      "loss": 0.57,
       "step": 251
     },
     {
       "epoch": 0.45528455284552843,
-      "grad_norm": 0.20579485783023402,
+      "grad_norm": 0.3272145932882665,
       "learning_rate": 1.8512589442218358e-05,
-      "loss": 0.5472,
+      "loss": 0.5523,
       "step": 252
     },
     {
       "epoch": 0.45709123757904246,
-      "grad_norm": 0.20450811714095307,
+      "grad_norm": 0.3089325909787182,
       "learning_rate": 1.8425263051659838e-05,
-      "loss": 0.5701,
+      "loss": 0.5751,
       "step": 253
     },
     {
       "epoch": 0.45889792231255644,
-      "grad_norm": 0.2795902989622695,
+      "grad_norm": 0.36927293242924614,
       "learning_rate": 1.8337814009344716e-05,
-      "loss": 0.5692,
+      "loss": 0.5736,
       "step": 254
     },
     {
       "epoch": 0.46070460704607047,
-      "grad_norm": 0.26043149982302555,
+      "grad_norm": 0.31333756675521174,
       "learning_rate": 1.8250245446646707e-05,
-      "loss": 0.5627,
+      "loss": 0.5667,
       "step": 255
     },
     {
       "epoch": 0.46251129177958444,
-      "grad_norm": 0.2115149201329313,
+      "grad_norm": 0.3194088697708899,
       "learning_rate": 1.8162560499219286e-05,
-      "loss": 0.563,
+      "loss": 0.567,
       "step": 256
     },
     {
       "epoch": 0.4643179765130985,
-      "grad_norm": 0.2863165296718786,
+      "grad_norm": 0.3358076681829595,
       "learning_rate": 1.807476230688346e-05,
-      "loss": 0.5726,
+      "loss": 0.5768,
       "step": 257
     },
     {
       "epoch": 0.46612466124661245,
-      "grad_norm": 0.22723135404479125,
+      "grad_norm": 0.3507131609472825,
       "learning_rate": 1.7986854013515274e-05,
-      "loss": 0.5682,
+      "loss": 0.573,
       "step": 258
     },
     {
       "epoch": 0.4679313459801265,
-      "grad_norm": 0.21269365868540174,
+      "grad_norm": 0.333864958674876,
       "learning_rate": 1.78988387669333e-05,
-      "loss": 0.5718,
+      "loss": 0.5759,
       "step": 259
     },
     {
       "epoch": 0.46973803071364045,
-      "grad_norm": 0.27706961689827647,
+      "grad_norm": 0.33675575200680075,
       "learning_rate": 1.781071971878587e-05,
-      "loss": 0.5852,
+      "loss": 0.5908,
       "step": 260
     },
     {
       "epoch": 0.4715447154471545,
-      "grad_norm": 0.21919053031721478,
+      "grad_norm": 0.32765202872639676,
       "learning_rate": 1.7722500024438244e-05,
-      "loss": 0.5766,
+      "loss": 0.5812,
       "step": 261
     },
     {
       "epoch": 0.47335140018066846,
-      "grad_norm": 0.2458759125768864,
+      "grad_norm": 0.36381544719415326,
       "learning_rate": 1.7634182842859628e-05,
-      "loss": 0.565,
+      "loss": 0.5693,
       "step": 262
     },
     {
       "epoch": 0.4751580849141825,
-      "grad_norm": 0.2489757112909537,
+      "grad_norm": 0.4851712003668751,
       "learning_rate": 1.7545771336510033e-05,
-      "loss": 0.5695,
+      "loss": 0.5736,
       "step": 263
     },
     {
       "epoch": 0.47696476964769646,
-      "grad_norm": 0.2101979687684982,
+      "grad_norm": 0.41291312521131196,
       "learning_rate": 1.7457268671227067e-05,
-      "loss": 0.5759,
+      "loss": 0.5784,
       "step": 264
     },
     {
       "epoch": 0.4787714543812105,
-      "grad_norm": 0.2500578266252668,
+      "grad_norm": 0.2955323853549614,
       "learning_rate": 1.736867801611254e-05,
-      "loss": 0.5569,
+      "loss": 0.5607,
       "step": 265
     },
     {
       "epoch": 0.48057813911472447,
-      "grad_norm": 0.19262415552045406,
+      "grad_norm": 0.37616325609332063,
       "learning_rate": 1.728000254341901e-05,
-      "loss": 0.5813,
+      "loss": 0.5861,
       "step": 266
     },
     {
       "epoch": 0.4823848238482385,
-      "grad_norm": 0.19478984136741012,
+      "grad_norm": 0.29381479811591366,
       "learning_rate": 1.7191245428436175e-05,
-      "loss": 0.5653,
+      "loss": 0.5699,
       "step": 267
     },
     {
       "epoch": 0.48419150858175247,
-      "grad_norm": 0.23871243123226799,
+      "grad_norm": 0.32495443011128766,
       "learning_rate": 1.7102409849377188e-05,
-      "loss": 0.5661,
+      "loss": 0.5697,
       "step": 268
     },
     {
       "epoch": 0.4859981933152665,
-      "grad_norm": 0.17665577180387385,
+      "grad_norm": 0.2624608390186021,
       "learning_rate": 1.7013498987264832e-05,
-      "loss": 0.5727,
+      "loss": 0.5769,
       "step": 269
     },
     {
       "epoch": 0.4878048780487805,
-      "grad_norm": 0.21091637968705407,
+      "grad_norm": 0.30513009228861504,
       "learning_rate": 1.6924516025817636e-05,
-      "loss": 0.5544,
+      "loss": 0.5586,
       "step": 270
     },
     {
       "epoch": 0.4896115627822945,
-      "grad_norm": 0.25253202754252524,
+      "grad_norm": 0.2900302075914806,
       "learning_rate": 1.683546415133584e-05,
-      "loss": 0.5761,
+      "loss": 0.5808,
       "step": 271
     },
     {
       "epoch": 0.4914182475158085,
-      "grad_norm": 0.2150474071339765,
+      "grad_norm": 0.30132501927469585,
       "learning_rate": 1.6746346552587342e-05,
-      "loss": 0.5704,
+      "loss": 0.5759,
       "step": 272
     },
     {
       "epoch": 0.4932249322493225,
-      "grad_norm": 0.24422291161795256,
+      "grad_norm": 0.3040014918275698,
       "learning_rate": 1.665716642069349e-05,
-      "loss": 0.5785,
+      "loss": 0.5823,
       "step": 273
     },
     {
       "epoch": 0.4950316169828365,
-      "grad_norm": 0.19465875800542706,
+      "grad_norm": 0.2581541065065194,
       "learning_rate": 1.6567926949014805e-05,
-      "loss": 0.5658,
+      "loss": 0.5705,
       "step": 274
     },
     {
       "epoch": 0.4968383017163505,
-      "grad_norm": 0.20059784342394038,
+      "grad_norm": 0.3506532731278841,
       "learning_rate": 1.6478631333036655e-05,
-      "loss": 0.5607,
+      "loss": 0.5654,
       "step": 275
     },
     {
       "epoch": 0.4986449864498645,
-      "grad_norm": 0.19428493006954836,
+      "grad_norm": 0.2841799020533544,
       "learning_rate": 1.638928277025482e-05,
-      "loss": 0.5523,
+      "loss": 0.5569,
       "step": 276
     },
     {
       "epoch": 0.5004516711833785,
-      "grad_norm": 0.190112158192132,
+      "grad_norm": 0.3147444372936614,
       "learning_rate": 1.6299884460061005e-05,
-      "loss": 0.5643,
+      "loss": 0.5683,
       "step": 277
     },
     {
       "epoch": 0.5022583559168925,
-      "grad_norm": 0.18317648770164868,
+      "grad_norm": 0.28903827689512285,
       "learning_rate": 1.621043960362826e-05,
-      "loss": 0.5509,
+      "loss": 0.5549,
       "step": 278
     },
     {
       "epoch": 0.5040650406504065,
-      "grad_norm": 0.20777757920037251,
+      "grad_norm": 0.31036698418346015,
       "learning_rate": 1.6120951403796367e-05,
-      "loss": 0.5659,
+      "loss": 0.5704,
       "step": 279
     },
     {
       "epoch": 0.5058717253839206,
-      "grad_norm": 0.24714269404949812,
+      "grad_norm": 0.31137986469945766,
       "learning_rate": 1.603142306495714e-05,
-      "loss": 0.5625,
+      "loss": 0.5661,
       "step": 280
     },
     {
       "epoch": 0.5076784101174345,
-      "grad_norm": 0.1771707758636966,
+      "grad_norm": 0.2971208151403228,
       "learning_rate": 1.5941857792939702e-05,
-      "loss": 0.5629,
+      "loss": 0.5673,
       "step": 281
     },
     {
       "epoch": 0.5094850948509485,
-      "grad_norm": 0.19209684745372535,
+      "grad_norm": 0.27859789351234376,
       "learning_rate": 1.585225879489567e-05,
-      "loss": 0.5577,
+      "loss": 0.562,
       "step": 282
     },
     {
       "epoch": 0.5112917795844625,
-      "grad_norm": 0.22597515112664288,
+      "grad_norm": 0.30768743394999065,
       "learning_rate": 1.5762629279184326e-05,
-      "loss": 0.5728,
+      "loss": 0.5779,
       "step": 283
     },
     {
       "epoch": 0.5130984643179766,
-      "grad_norm": 0.17250544743705307,
+      "grad_norm": 0.29234292940585527,
       "learning_rate": 1.5672972455257726e-05,
-      "loss": 0.5611,
+      "loss": 0.5662,
       "step": 284
     },
     {
       "epoch": 0.5149051490514905,
-      "grad_norm": 0.20410988336988065,
+      "grad_norm": 0.32303845185987906,
       "learning_rate": 1.5583291533545775e-05,
-      "loss": 0.5557,
+      "loss": 0.5603,
       "step": 285
     },
     {
       "epoch": 0.5167118337850045,
-      "grad_norm": 0.20391700905101098,
+      "grad_norm": 0.2994101822443779,
       "learning_rate": 1.549358972534128e-05,
-      "loss": 0.5538,
+      "loss": 0.5577,
       "step": 286
     },
     {
       "epoch": 0.5185185185185185,
-      "grad_norm": 0.18259503875294242,
+      "grad_norm": 0.3225727799664951,
       "learning_rate": 1.5403870242684942e-05,
-      "loss": 0.5754,
+      "loss": 0.5794,
       "step": 287
     },
     {
       "epoch": 0.5203252032520326,
-      "grad_norm": 0.18696008552172258,
+      "grad_norm": 0.2923618155390653,
       "learning_rate": 1.5314136298250355e-05,
-      "loss": 0.5723,
+      "loss": 0.577,
       "step": 288
     },
     {
       "epoch": 0.5221318879855466,
-      "grad_norm": 0.19653948540679103,
+      "grad_norm": 0.29720244961677544,
       "learning_rate": 1.5224391105228956e-05,
-      "loss": 0.5565,
+      "loss": 0.5603,
       "step": 289
     },
     {
       "epoch": 0.5239385727190605,
-      "grad_norm": 0.17253613754365116,
+      "grad_norm": 0.2989180878702361,
       "learning_rate": 1.5134637877214968e-05,
-      "loss": 0.5682,
+      "loss": 0.5727,
       "step": 290
     },
     {
       "epoch": 0.5257452574525745,
-      "grad_norm": 0.2056061343083173,
+      "grad_norm": 0.25519500575209336,
       "learning_rate": 1.5044879828090346e-05,
-      "loss": 0.5727,
+      "loss": 0.5773,
       "step": 291
     },
     {
       "epoch": 0.5275519421860885,
-      "grad_norm": 0.19894424925026757,
+      "grad_norm": 0.2830864583495515,
       "learning_rate": 1.4955120171909658e-05,
-      "loss": 0.5763,
+      "loss": 0.5816,
       "step": 292
     },
     {
       "epoch": 0.5293586269196026,
-      "grad_norm": 0.1659275906965013,
+      "grad_norm": 0.2764753853248595,
       "learning_rate": 1.4865362122785031e-05,
-      "loss": 0.5689,
+      "loss": 0.5744,
       "step": 293
     },
     {
       "epoch": 0.5311653116531165,
-      "grad_norm": 0.17986609793129524,
+      "grad_norm": 0.3547980154145159,
       "learning_rate": 1.4775608894771048e-05,
-      "loss": 0.5612,
+      "loss": 0.5664,
       "step": 294
     },
     {
       "epoch": 0.5329719963866305,
-      "grad_norm": 0.232250278692599,
+      "grad_norm": 0.29803300652072573,
       "learning_rate": 1.4685863701749648e-05,
-      "loss": 0.5589,
+      "loss": 0.5639,
       "step": 295
     },
     {
       "epoch": 0.5347786811201445,
-      "grad_norm": 0.22464497097731725,
+      "grad_norm": 0.290491851084967,
       "learning_rate": 1.4596129757315062e-05,
-      "loss": 0.5628,
+      "loss": 0.5681,
       "step": 296
     },
     {
       "epoch": 0.5365853658536586,
-      "grad_norm": 0.1952420747386217,
+      "grad_norm": 0.3389739871118521,
       "learning_rate": 1.4506410274658718e-05,
-      "loss": 0.5682,
+      "loss": 0.5726,
       "step": 297
     },
     {
       "epoch": 0.5383920505871725,
-      "grad_norm": 0.2106666884744506,
+      "grad_norm": 0.3246943570187164,
       "learning_rate": 1.441670846645423e-05,
-      "loss": 0.5627,
+      "loss": 0.5676,
       "step": 298
     },
     {
       "epoch": 0.5401987353206865,
-      "grad_norm": 0.21061181918065086,
+      "grad_norm": 0.3242408286020586,
       "learning_rate": 1.4327027544742281e-05,
-      "loss": 0.5783,
+      "loss": 0.5839,
       "step": 299
     },
     {
       "epoch": 0.5420054200542005,
-      "grad_norm": 0.1861290798519661,
+      "grad_norm": 0.28501826222047294,
       "learning_rate": 1.4237370720815675e-05,
-      "loss": 0.5664,
+      "loss": 0.5717,
       "step": 300
     },
     {
       "epoch": 0.5438121047877146,
-      "grad_norm": 0.18862693026886343,
+      "grad_norm": 0.2801125597531159,
       "learning_rate": 1.4147741205104336e-05,
-      "loss": 0.5683,
+      "loss": 0.5742,
       "step": 301
     },
     {
       "epoch": 0.5456187895212286,
-      "grad_norm": 0.18676460565472633,
+      "grad_norm": 0.2865662680741476,
       "learning_rate": 1.40581422070603e-05,
-      "loss": 0.5796,
+      "loss": 0.5859,
       "step": 302
     },
     {
       "epoch": 0.5474254742547425,
-      "grad_norm": 0.1752400383590249,
+      "grad_norm": 0.2880283177460633,
       "learning_rate": 1.3968576935042864e-05,
-      "loss": 0.5608,
+      "loss": 0.5666,
       "step": 303
     },
     {
       "epoch": 0.5492321589882565,
-      "grad_norm": 0.18836688433228918,
+      "grad_norm": 0.29776279014261486,
       "learning_rate": 1.3879048596203637e-05,
-      "loss": 0.5561,
+      "loss": 0.5614,
       "step": 304
     },
     {
       "epoch": 0.5510388437217706,
-      "grad_norm": 0.19530506408089982,
+      "grad_norm": 0.3046214139843183,
       "learning_rate": 1.3789560396371742e-05,
-      "loss": 0.5622,
+      "loss": 0.5679,
       "step": 305
     },
     {
       "epoch": 0.5528455284552846,
-      "grad_norm": 0.17965863699921675,
+      "grad_norm": 0.31091017665288007,
       "learning_rate": 1.3700115539938995e-05,
-      "loss": 0.5605,
+      "loss": 0.5652,
       "step": 306
     },
     {
       "epoch": 0.5546522131887985,
-      "grad_norm": 0.21750349150527712,
+      "grad_norm": 0.3342055750377691,
       "learning_rate": 1.3610717229745183e-05,
-      "loss": 0.578,
+      "loss": 0.5831,
       "step": 307
     },
     {
       "epoch": 0.5564588979223125,
-      "grad_norm": 0.19147158863524186,
+      "grad_norm": 0.2825943487551394,
       "learning_rate": 1.3521368666963348e-05,
-      "loss": 0.5572,
+      "loss": 0.5622,
       "step": 308
     },
     {
       "epoch": 0.5582655826558266,
-      "grad_norm": 0.18608236035974007,
+      "grad_norm": 0.42463521342108723,
       "learning_rate": 1.3432073050985201e-05,
-      "loss": 0.5669,
+      "loss": 0.5719,
       "step": 309
     },
     {
       "epoch": 0.5600722673893406,
-      "grad_norm": 0.18440593641158914,
+      "grad_norm": 0.24690027892733366,
       "learning_rate": 1.3342833579306511e-05,
-      "loss": 0.5492,
+      "loss": 0.5537,
       "step": 310
     },
     {
       "epoch": 0.5618789521228545,
-      "grad_norm": 0.1843799787511743,
+      "grad_norm": 0.2782245278873188,
       "learning_rate": 1.3253653447412658e-05,
-      "loss": 0.5685,
+      "loss": 0.5738,
       "step": 311
     },
     {
       "epoch": 0.5636856368563685,
-      "grad_norm": 0.20261725696919525,
+      "grad_norm": 0.269514288271004,
       "learning_rate": 1.3164535848664167e-05,
-      "loss": 0.5674,
+      "loss": 0.5728,
       "step": 312
     },
     {
       "epoch": 0.5654923215898826,
-      "grad_norm": 0.17558184672847507,
+      "grad_norm": 0.28864669964933554,
       "learning_rate": 1.3075483974182368e-05,
-      "loss": 0.5684,
+      "loss": 0.5731,
       "step": 313
     },
     {
       "epoch": 0.5672990063233966,
-      "grad_norm": 0.17907331587608558,
+      "grad_norm": 0.2612891164500954,
       "learning_rate": 1.2986501012735174e-05,
-      "loss": 0.5609,
+      "loss": 0.5663,
       "step": 314
     },
     {
       "epoch": 0.5691056910569106,
-      "grad_norm": 0.1934736256775333,
+      "grad_norm": 0.2551741665283283,
       "learning_rate": 1.2897590150622812e-05,
-      "loss": 0.5613,
+      "loss": 0.5662,
       "step": 315
     },
     {
       "epoch": 0.5709123757904245,
-      "grad_norm": 0.17567517580021857,
+      "grad_norm": 0.28291261096839593,
       "learning_rate": 1.2808754571563827e-05,
-      "loss": 0.5573,
+      "loss": 0.5623,
       "step": 316
     },
     {
       "epoch": 0.5727190605239386,
-      "grad_norm": 0.19690748080324788,
+      "grad_norm": 0.33675663529353467,
       "learning_rate": 1.2719997456580993e-05,
-      "loss": 0.5525,
+      "loss": 0.5573,
       "step": 317
     },
     {
       "epoch": 0.5745257452574526,
-      "grad_norm": 0.20212907631156332,
+      "grad_norm": 0.242227937824476,
       "learning_rate": 1.2631321983887464e-05,
-      "loss": 0.5654,
+      "loss": 0.5701,
       "step": 318
     },
     {
       "epoch": 0.5763324299909666,
-      "grad_norm": 0.19608541723414,
+      "grad_norm": 1.7361889228979208,
       "learning_rate": 1.2542731328772936e-05,
-      "loss": 0.5744,
+      "loss": 0.5796,
       "step": 319
     },
     {
       "epoch": 0.5781391147244805,
-      "grad_norm": 0.20857428025225674,
+      "grad_norm": 0.29371094872544395,
       "learning_rate": 1.2454228663489967e-05,
-      "loss": 0.5569,
+      "loss": 0.5622,
       "step": 320
     },
     {
       "epoch": 0.5799457994579946,
-      "grad_norm": 0.17609191697173338,
+      "grad_norm": 0.275215385186721,
       "learning_rate": 1.2365817157140373e-05,
-      "loss": 0.5541,
+      "loss": 0.5595,
       "step": 321
     },
     {
       "epoch": 0.5817524841915086,
-      "grad_norm": 0.18308421719879198,
+      "grad_norm": 0.2925560707999663,
       "learning_rate": 1.2277499975561755e-05,
-      "loss": 0.5753,
+      "loss": 0.5807,
       "step": 322
     },
     {
       "epoch": 0.5835591689250226,
-      "grad_norm": 0.17484046660309832,
+      "grad_norm": 0.4068710625406492,
       "learning_rate": 1.2189280281214128e-05,
-      "loss": 0.5609,
+      "loss": 0.566,
       "step": 323
     },
     {
       "epoch": 0.5853658536585366,
-      "grad_norm": 0.17252551757672757,
+      "grad_norm": 0.2628331351797529,
       "learning_rate": 1.2101161233066704e-05,
-      "loss": 0.5419,
+      "loss": 0.5469,
       "step": 324
     },
     {
       "epoch": 0.5871725383920506,
-      "grad_norm": 0.17553463607522976,
+      "grad_norm": 0.2911285026218736,
       "learning_rate": 1.201314598648473e-05,
-      "loss": 0.5591,
+      "loss": 0.565,
       "step": 325
     },
     {
       "epoch": 0.5889792231255646,
-      "grad_norm": 0.17326333720847636,
+      "grad_norm": 0.2831006227507552,
       "learning_rate": 1.1925237693116546e-05,
-      "loss": 0.5689,
+      "loss": 0.5747,
       "step": 326
     },
     {
       "epoch": 0.5907859078590786,
-      "grad_norm": 0.18687445963344007,
+      "grad_norm": 0.26598868778075935,
       "learning_rate": 1.1837439500780718e-05,
-      "loss": 0.5553,
+      "loss": 0.5602,
       "step": 327
     },
     {
       "epoch": 0.5925925925925926,
-      "grad_norm": 0.19428514170213504,
+      "grad_norm": 0.3558168004733384,
       "learning_rate": 1.1749754553353297e-05,
-      "loss": 0.562,
+      "loss": 0.5676,
       "step": 328
     },
     {
       "epoch": 0.5943992773261066,
-      "grad_norm": 0.17110021647920834,
+      "grad_norm": 0.25067605853110625,
       "learning_rate": 1.1662185990655285e-05,
-      "loss": 0.5518,
+      "loss": 0.5569,
       "step": 329
     },
     {
       "epoch": 0.5962059620596206,
-      "grad_norm": 0.18984185404031334,
+      "grad_norm": 0.3465570264289409,
       "learning_rate": 1.1574736948340163e-05,
-      "loss": 0.559,
+      "loss": 0.5644,
       "step": 330
     },
     {
       "epoch": 0.5980126467931346,
-      "grad_norm": 0.16874712774452827,
+      "grad_norm": 0.2555984652347325,
       "learning_rate": 1.1487410557781642e-05,
-      "loss": 0.5632,
+      "loss": 0.5689,
       "step": 331
     },
     {
       "epoch": 0.5998193315266486,
-      "grad_norm": 0.1806801244410073,
+      "grad_norm": 1.3045885103794765,
       "learning_rate": 1.140020994596151e-05,
-      "loss": 0.5467,
+      "loss": 0.5525,
       "step": 332
     },
     {
       "epoch": 0.6016260162601627,
-      "grad_norm": 0.17626693416571312,
+      "grad_norm": 0.269714982616486,
       "learning_rate": 1.1313138235357683e-05,
-      "loss": 0.5674,
+      "loss": 0.5735,
       "step": 333
     },
     {
       "epoch": 0.6034327009936766,
-      "grad_norm": 0.1770136401046488,
+      "grad_norm": 0.2468409802180684,
       "learning_rate": 1.1226198543832373e-05,
-      "loss": 0.5798,
+      "loss": 0.5864,
       "step": 334
     },
     {
       "epoch": 0.6052393857271906,
-      "grad_norm": 0.1910726991764081,
+      "grad_norm": 0.2359964626102903,
       "learning_rate": 1.1139393984520467e-05,
-      "loss": 0.563,
+      "loss": 0.5688,
       "step": 335
     },
     {
       "epoch": 0.6070460704607046,
-      "grad_norm": 0.1718176664612323,
+      "grad_norm": 0.2792218994212215,
       "learning_rate": 1.1052727665718027e-05,
-      "loss": 0.5596,
+      "loss": 0.5651,
       "step": 336
     },
     {
       "epoch": 0.6088527551942186,
-      "grad_norm": 0.17250346288280585,
+      "grad_norm": 0.24660150003449235,
       "learning_rate": 1.0966202690771015e-05,
-      "loss": 0.5611,
+      "loss": 0.5673,
       "step": 337
     },
     {
       "epoch": 0.6106594399277326,
-      "grad_norm": 0.19111207535840807,
+      "grad_norm": 0.28553688739699007,
       "learning_rate": 1.087982215796415e-05,
-      "loss": 0.5572,
+      "loss": 0.5628,
       "step": 338
     },
     {
       "epoch": 0.6124661246612466,
-      "grad_norm": 0.1686795998299278,
+      "grad_norm": 0.24555542198671096,
       "learning_rate": 1.079358916040996e-05,
-      "loss": 0.5657,
+      "loss": 0.5717,
       "step": 339
     },
     {
       "epoch": 0.6142728093947606,
-      "grad_norm": 0.16862020455628596,
+      "grad_norm": 0.25213878756115177,
       "learning_rate": 1.0707506785938052e-05,
-      "loss": 0.5611,
+      "loss": 0.567,
       "step": 340
     },
     {
       "epoch": 0.6160794941282746,
-      "grad_norm": 0.1637656583385122,
+      "grad_norm": 0.2316189652392717,
       "learning_rate": 1.0621578116984501e-05,
-      "loss": 0.5593,
+      "loss": 0.5649,
       "step": 341
     },
     {
       "epoch": 0.6178861788617886,
-      "grad_norm": 0.16170288648566758,
+      "grad_norm": 0.2642230456418394,
       "learning_rate": 1.0535806230481525e-05,
-      "loss": 0.5628,
+      "loss": 0.5682,
       "step": 342
     },
     {
       "epoch": 0.6196928635953026,
-      "grad_norm": 0.1717924247456263,
+      "grad_norm": 0.25607403892541075,
       "learning_rate": 1.0450194197747252e-05,
-      "loss": 0.5567,
+      "loss": 0.5626,
       "step": 343
     },
     {
       "epoch": 0.6214995483288166,
-      "grad_norm": 0.16174500558159677,
+      "grad_norm": 0.2454366513142169,
       "learning_rate": 1.036474508437579e-05,
-      "loss": 0.5464,
+      "loss": 0.5524,
       "step": 344
     },
     {
       "epoch": 0.6233062330623306,
-      "grad_norm": 0.16381172759058904,
+      "grad_norm": 0.25692513164041986,
       "learning_rate": 1.0279461950127422e-05,
-      "loss": 0.5608,
+      "loss": 0.5666,
       "step": 345
     },
     {
       "epoch": 0.6251129177958447,
-      "grad_norm": 0.17001972937010654,
+      "grad_norm": 0.2567223339593557,
       "learning_rate": 1.0194347848819055e-05,
-      "loss": 0.5623,
+      "loss": 0.5675,
       "step": 346
     },
     {
       "epoch": 0.6269196025293586,
-      "grad_norm": 0.17303311757989492,
+      "grad_norm": 0.24303756149752093,
       "learning_rate": 1.0109405828214872e-05,
-      "loss": 0.5482,
+      "loss": 0.5534,
       "step": 347
     },
     {
       "epoch": 0.6287262872628726,
-      "grad_norm": 0.1882537905657933,
+      "grad_norm": 0.3469962933198629,
       "learning_rate": 1.0024638929917192e-05,
-      "loss": 0.5492,
+      "loss": 0.5543,
       "step": 348
     },
     {
       "epoch": 0.6305329719963866,
-      "grad_norm": 0.15457360070696674,
+      "grad_norm": 0.29482139739377156,
       "learning_rate": 9.940050189257552e-06,
-      "loss": 0.5521,
+      "loss": 0.5582,
       "step": 349
     },
     {
       "epoch": 0.6323396567299007,
-      "grad_norm": 0.15615688544922562,
+      "grad_norm": 0.2543061317007475,
       "learning_rate": 9.85564263518803e-06,
-      "loss": 0.5643,
+      "loss": 0.5703,
       "step": 350
     },
     {
       "epoch": 0.6341463414634146,
-      "grad_norm": 0.16319062879774318,
+      "grad_norm": 0.2598563996097049,
       "learning_rate": 9.771419290172776e-06,
-      "loss": 0.554,
+      "loss": 0.5594,
       "step": 351
     },
     {
       "epoch": 0.6359530261969286,
-      "grad_norm": 0.15770060432248675,
+      "grad_norm": 0.28480208198026236,
       "learning_rate": 9.68738317007978e-06,
-      "loss": 0.5682,
+      "loss": 0.5738,
       "step": 352
     },
     {
       "epoch": 0.6377597109304426,
-      "grad_norm": 0.16898427412267347,
+      "grad_norm": 0.24709090613155346,
       "learning_rate": 9.6035372840729e-06,
-      "loss": 0.568,
+      "loss": 0.5738,
       "step": 353
     },
     {
       "epoch": 0.6395663956639567,
-      "grad_norm": 0.15918955141776597,
+      "grad_norm": 0.2671031729169573,
       "learning_rate": 9.519884634504074e-06,
-      "loss": 0.5631,
+      "loss": 0.5691,
       "step": 354
     },
     {
       "epoch": 0.6413730803974707,
-      "grad_norm": 0.1635896927311437,
+      "grad_norm": 0.22856827107644548,
       "learning_rate": 9.436428216805862e-06,
-      "loss": 0.5524,
+      "loss": 0.5579,
       "step": 355
     },
     {
       "epoch": 0.6431797651309846,
-      "grad_norm": 0.15329421947992738,
+      "grad_norm": 0.2544739414287776,
       "learning_rate": 9.35317101938412e-06,
-      "loss": 0.566,
+      "loss": 0.5723,
       "step": 356
     },
     {
       "epoch": 0.6449864498644986,
-      "grad_norm": 0.18290183186305234,
+      "grad_norm": 0.25959041144887784,
       "learning_rate": 9.270116023511063e-06,
-      "loss": 0.5676,
+      "loss": 0.5736,
       "step": 357
     },
     {
       "epoch": 0.6467931345980127,
-      "grad_norm": 0.18253658487556157,
+      "grad_norm": 0.4509629870171336,
       "learning_rate": 9.187266203218457e-06,
-      "loss": 0.5729,
+      "loss": 0.5788,
       "step": 358
     },
     {
       "epoch": 0.6485998193315267,
-      "grad_norm": 0.1791341463101353,
+      "grad_norm": 0.2554658264431559,
       "learning_rate": 9.104624525191147e-06,
-      "loss": 0.546,
+      "loss": 0.5516,
       "step": 359
     },
     {
       "epoch": 0.6504065040650406,
-      "grad_norm": 0.14674620232896526,
+      "grad_norm": 0.22730131545870966,
       "learning_rate": 9.022193948660825e-06,
-      "loss": 0.5548,
+      "loss": 0.5601,
       "step": 360
     },
     {
       "epoch": 0.6522131887985546,
-      "grad_norm": 0.19660949623094373,
+      "grad_norm": 0.2705934074611756,
       "learning_rate": 8.939977425300072e-06,
-      "loss": 0.56,
+      "loss": 0.5657,
       "step": 361
     },
     {
       "epoch": 0.6540198735320687,
-      "grad_norm": 0.1876501178479123,
+      "grad_norm": 0.23508126204947832,
       "learning_rate": 8.857977899116645e-06,
-      "loss": 0.5717,
+      "loss": 0.5772,
       "step": 362
     },
     {
       "epoch": 0.6558265582655827,
-      "grad_norm": 0.21507796351172995,
+      "grad_norm": 0.2670484472522681,
       "learning_rate": 8.776198306348072e-06,
-      "loss": 0.5618,
+      "loss": 0.5679,
       "step": 363
     },
     {
       "epoch": 0.6576332429990966,
-      "grad_norm": 0.18228320215554383,
+      "grad_norm": 0.25807775800350113,
       "learning_rate": 8.69464157535652e-06,
-      "loss": 0.5612,
+      "loss": 0.5678,
       "step": 364
     },
     {
       "epoch": 0.6594399277326106,
-      "grad_norm": 0.1757206345144353,
+      "grad_norm": 0.4143373428623339,
       "learning_rate": 8.61331062652391e-06,
-      "loss": 0.5534,
+      "loss": 0.5593,
       "step": 365
     },
     {
       "epoch": 0.6612466124661247,
-      "grad_norm": 0.19026770352731234,
+      "grad_norm": 0.25432160506360196,
       "learning_rate": 8.532208372147376e-06,
-      "loss": 0.5746,
+      "loss": 0.5811,
       "step": 366
     },
     {
       "epoch": 0.6630532971996387,
-      "grad_norm": 0.1865886995183675,
+      "grad_norm": 0.2218223608724009,
       "learning_rate": 8.451337716334953e-06,
-      "loss": 0.5556,
+      "loss": 0.5612,
       "step": 367
     },
     {
       "epoch": 0.6648599819331527,
-      "grad_norm": 0.1797695927669861,
+      "grad_norm": 0.26338098583116515,
       "learning_rate": 8.370701554901603e-06,
-      "loss": 0.5675,
+      "loss": 0.573,
       "step": 368
     },
     {
       "epoch": 0.6666666666666666,
-      "grad_norm": 0.17670482413719021,
+      "grad_norm": 0.27305656638372205,
       "learning_rate": 8.290302775265509e-06,
-      "loss": 0.5686,
+      "loss": 0.5743,
       "step": 369
     },
     {
       "epoch": 0.6684733514001807,
-      "grad_norm": 0.1701577263486839,
+      "grad_norm": 0.28354813004240265,
       "learning_rate": 8.210144256344714e-06,
-      "loss": 0.5668,
+      "loss": 0.5732,
       "step": 370
     },
     {
       "epoch": 0.6702800361336947,
-      "grad_norm": 0.2061129978805257,
+      "grad_norm": 0.2540985448161381,
       "learning_rate": 8.130228868453989e-06,
-      "loss": 0.555,
+      "loss": 0.5604,
       "step": 371
     },
     {
       "epoch": 0.6720867208672087,
-      "grad_norm": 0.1969046824239301,
+      "grad_norm": 0.23334255114709357,
       "learning_rate": 8.050559473202078e-06,
-      "loss": 0.5727,
+      "loss": 0.578,
       "step": 372
     },
     {
       "epoch": 0.6738934056007226,
-      "grad_norm": 0.15311615793545202,
+      "grad_norm": 0.23245111955419623,
       "learning_rate": 7.971138923389221e-06,
-      "loss": 0.5583,
+      "loss": 0.5634,
       "step": 373
     },
     {
       "epoch": 0.6757000903342367,
-      "grad_norm": 0.18976350186526125,
+      "grad_norm": 0.33850041377146883,
       "learning_rate": 7.89197006290502e-06,
-      "loss": 0.5473,
+      "loss": 0.5532,
       "step": 374
     },
     {
       "epoch": 0.6775067750677507,
-      "grad_norm": 0.1879352111528028,
+      "grad_norm": 0.24306579520579513,
       "learning_rate": 7.813055726626579e-06,
-      "loss": 0.5619,
+      "loss": 0.567,
       "step": 375
     },
     {
       "epoch": 0.6793134598012647,
-      "grad_norm": 0.20499323771840253,
+      "grad_norm": 0.23084325043799506,
       "learning_rate": 7.734398740316999e-06,
-      "loss": 0.5467,
+      "loss": 0.552,
       "step": 376
     },
     {
       "epoch": 0.6811201445347786,
-      "grad_norm": 0.17137028091933101,
+      "grad_norm": 0.24330485181839703,
       "learning_rate": 7.65600192052422e-06,
-      "loss": 0.563,
+      "loss": 0.5686,
       "step": 377
     },
     {
       "epoch": 0.6829268292682927,
-      "grad_norm": 0.18099761803792097,
+      "grad_norm": 0.2537299633849379,
       "learning_rate": 7.577868074480107e-06,
-      "loss": 0.5634,
+      "loss": 0.5695,
       "step": 378
     },
     {
       "epoch": 0.6847335140018067,
-      "grad_norm": 0.19896866364484633,
+      "grad_norm": 0.2324650195741195,
       "learning_rate": 7.500000000000004e-06,
-      "loss": 0.5664,
+      "loss": 0.5723,
       "step": 379
     },
     {
       "epoch": 0.6865401987353207,
-      "grad_norm": 0.1884102435991128,
+      "grad_norm": 0.2559577756429273,
       "learning_rate": 7.422400485382484e-06,
-      "loss": 0.5473,
+      "loss": 0.553,
       "step": 380
     },
     {
       "epoch": 0.6883468834688347,
-      "grad_norm": 0.17513345385135934,
+      "grad_norm": 0.25311375180891915,
       "learning_rate": 7.345072309309557e-06,
-      "loss": 0.5638,
+      "loss": 0.5694,
       "step": 381
     },
     {
       "epoch": 0.6901535682023487,
-      "grad_norm": 0.1733806383727063,
+      "grad_norm": 0.2344778183121317,
       "learning_rate": 7.268018240747129e-06,
-      "loss": 0.5515,
+      "loss": 0.5564,
       "step": 382
     },
     {
       "epoch": 0.6919602529358627,
-      "grad_norm": 0.18213237487884423,
+      "grad_norm": 0.21551852417528822,
       "learning_rate": 7.191241038845876e-06,
-      "loss": 0.5704,
+      "loss": 0.5767,
       "step": 383
     },
     {
       "epoch": 0.6937669376693767,
-      "grad_norm": 0.20160376104741037,
+      "grad_norm": 0.2715893346415364,
       "learning_rate": 7.114743452842428e-06,
-      "loss": 0.5681,
+      "loss": 0.5738,
       "step": 384
     },
     {
       "epoch": 0.6955736224028907,
-      "grad_norm": 0.14133588164420194,
+      "grad_norm": 0.23073802447177677,
       "learning_rate": 7.038528221960959e-06,
-      "loss": 0.5504,
+      "loss": 0.5558,
       "step": 385
     },
     {
       "epoch": 0.6973803071364046,
-      "grad_norm": 0.1554959410199801,
+      "grad_norm": 0.240908436382785,
       "learning_rate": 6.962598075315047e-06,
-      "loss": 0.5622,
+      "loss": 0.5677,
       "step": 386
     },
     {
       "epoch": 0.6991869918699187,
-      "grad_norm": 0.17804707625227728,
+      "grad_norm": 0.256554182250642,
       "learning_rate": 6.886955731810009e-06,
-      "loss": 0.5622,
+      "loss": 0.5679,
       "step": 387
     },
     {
       "epoch": 0.7009936766034327,
-      "grad_norm": 0.1770700659186734,
+      "grad_norm": 0.39527336236612054,
       "learning_rate": 6.811603900045477e-06,
-      "loss": 0.5612,
+      "loss": 0.5663,
       "step": 388
     },
     {
       "epoch": 0.7028003613369467,
-      "grad_norm": 0.15981371678518486,
+      "grad_norm": 0.25185271897899014,
       "learning_rate": 6.736545278218464e-06,
-      "loss": 0.5642,
+      "loss": 0.5698,
       "step": 389
     },
     {
       "epoch": 0.7046070460704607,
-      "grad_norm": 0.1530064162827602,
+      "grad_norm": 0.3166682567846798,
       "learning_rate": 6.661782554026735e-06,
-      "loss": 0.5584,
+      "loss": 0.5638,
       "step": 390
     },
     {
       "epoch": 0.7064137308039747,
-      "grad_norm": 0.16596305105845469,
+      "grad_norm": 0.2319838260541366,
       "learning_rate": 6.587318404572537e-06,
-      "loss": 0.5583,
+      "loss": 0.564,
       "step": 391
     },
     {
       "epoch": 0.7082204155374887,
-      "grad_norm": 0.16327470933567317,
+      "grad_norm": 0.22505306970140945,
       "learning_rate": 6.513155496266783e-06,
-      "loss": 0.5641,
+      "loss": 0.5703,
       "step": 392
     },
     {
       "epoch": 0.7100271002710027,
-      "grad_norm": 0.1574337263337296,
+      "grad_norm": 0.24527359143209262,
       "learning_rate": 6.439296484733526e-06,
-      "loss": 0.5689,
+      "loss": 0.575,
       "step": 393
     },
     {
       "epoch": 0.7118337850045167,
-      "grad_norm": 0.16018951967211098,
+      "grad_norm": 0.21842675068411357,
       "learning_rate": 6.3657440147149e-06,
-      "loss": 0.5716,
+      "loss": 0.5776,
       "step": 394
     },
     {
       "epoch": 0.7136404697380307,
-      "grad_norm": 0.14974112027258443,
+      "grad_norm": 0.22526457972223854,
       "learning_rate": 6.292500719976397e-06,
-      "loss": 0.5659,
+      "loss": 0.5722,
       "step": 395
     },
     {
       "epoch": 0.7154471544715447,
-      "grad_norm": 0.16158709365351387,
+      "grad_norm": 0.2386933725648013,
       "learning_rate": 6.219569223212585e-06,
-      "loss": 0.56,
+      "loss": 0.5659,
       "step": 396
     },
     {
       "epoch": 0.7172538392050587,
-      "grad_norm": 0.15600729659089102,
+      "grad_norm": 0.2797910217583912,
       "learning_rate": 6.146952135953158e-06,
-      "loss": 0.5598,
+      "loss": 0.5663,
       "step": 397
     },
     {
       "epoch": 0.7190605239385727,
-      "grad_norm": 0.15620491253414515,
+      "grad_norm": 0.222440302196531,
       "learning_rate": 6.074652058469447e-06,
-      "loss": 0.5616,
+      "loss": 0.5677,
       "step": 398
     },
     {
       "epoch": 0.7208672086720868,
-      "grad_norm": 0.16380045583881614,
+      "grad_norm": 0.2244823119289259,
       "learning_rate": 6.0026715796812945e-06,
-      "loss": 0.5644,
+      "loss": 0.5694,
       "step": 399
     },
     {
       "epoch": 0.7226738934056007,
-      "grad_norm": 0.14967907113452944,
+      "grad_norm": 0.22028730249011422,
       "learning_rate": 5.931013277064377e-06,
-      "loss": 0.5609,
+      "loss": 0.567,
       "step": 400
     },
     {
       "epoch": 0.7244805781391147,
-      "grad_norm": 0.15280123889646421,
+      "grad_norm": 0.21677756210322688,
       "learning_rate": 5.859679716557879e-06,
-      "loss": 0.5626,
+      "loss": 0.5688,
       "step": 401
     },
     {
       "epoch": 0.7262872628726287,
-      "grad_norm": 0.16341975790372415,
+      "grad_norm": 0.4344517020406158,
       "learning_rate": 5.788673452472622e-06,
-      "loss": 0.5696,
+      "loss": 0.5752,
       "step": 402
     },
     {
       "epoch": 0.7280939476061428,
-      "grad_norm": 0.1562147019531761,
+      "grad_norm": 0.24869038865191093,
       "learning_rate": 5.717997027399624e-06,
-      "loss": 0.5516,
+      "loss": 0.5567,
       "step": 403
     },
     {
       "epoch": 0.7299006323396567,
-      "grad_norm": 0.15660229823809,
+      "grad_norm": 0.23064985584415565,
       "learning_rate": 5.647652972118998e-06,
-      "loss": 0.5569,
+      "loss": 0.5626,
       "step": 404
     },
     {
       "epoch": 0.7317073170731707,
-      "grad_norm": 0.16555666302974795,
+      "grad_norm": 0.22343320849286405,
       "learning_rate": 5.577643805509404e-06,
-      "loss": 0.5549,
+      "loss": 0.5606,
       "step": 405
     },
     {
       "epoch": 0.7335140018066847,
-      "grad_norm": 0.1507767089868411,
+      "grad_norm": 0.2410416469241811,
       "learning_rate": 5.507972034457786e-06,
-      "loss": 0.5615,
+      "loss": 0.567,
       "step": 406
     },
     {
       "epoch": 0.7353206865401988,
-      "grad_norm": 0.1557734692497816,
+      "grad_norm": 0.21682033645252197,
       "learning_rate": 5.438640153769654e-06,
-      "loss": 0.5534,
+      "loss": 0.5594,
       "step": 407
     },
     {
       "epoch": 0.7371273712737128,
-      "grad_norm": 0.14782949914277185,
+      "grad_norm": 0.21362885977787827,
       "learning_rate": 5.369650646079719e-06,
-      "loss": 0.547,
+      "loss": 0.5527,
       "step": 408
     },
     {
       "epoch": 0.7389340560072267,
-      "grad_norm": 0.17200870547934033,
+      "grad_norm": 0.2501581124737507,
       "learning_rate": 5.301005981763007e-06,
-      "loss": 0.5669,
+      "loss": 0.5723,
       "step": 409
     },
     {
       "epoch": 0.7407407407407407,
-      "grad_norm": 0.16956590081345058,
+      "grad_norm": 0.21950559446814377,
       "learning_rate": 5.2327086188463964e-06,
-      "loss": 0.5672,
+      "loss": 0.5733,
       "step": 410
     },
     {
       "epoch": 0.7425474254742548,
-      "grad_norm": 0.17288972695084148,
+      "grad_norm": 0.22874967519790235,
       "learning_rate": 5.164761002920615e-06,
-      "loss": 0.5695,
+      "loss": 0.5757,
       "step": 411
     },
     {
       "epoch": 0.7443541102077688,
-      "grad_norm": 0.15302551358387134,
+      "grad_norm": 0.23529791112541182,
       "learning_rate": 5.097165567052646e-06,
-      "loss": 0.543,
+      "loss": 0.5478,
       "step": 412
     },
     {
       "epoch": 0.7461607949412827,
-      "grad_norm": 0.1659332169274283,
+      "grad_norm": 0.24186795740106914,
       "learning_rate": 5.029924731698619e-06,
-      "loss": 0.5586,
+      "loss": 0.5641,
       "step": 413
     },
     {
       "epoch": 0.7479674796747967,
-      "grad_norm": 0.14771275975107256,
+      "grad_norm": 0.22241278304703085,
       "learning_rate": 4.963040904617131e-06,
-      "loss": 0.5466,
+      "loss": 0.5528,
       "step": 414
     },
     {
       "epoch": 0.7497741644083108,
-      "grad_norm": 0.14811614613987537,
+      "grad_norm": 0.23230777162136979,
       "learning_rate": 4.8965164807830355e-06,
-      "loss": 0.568,
+      "loss": 0.5735,
       "step": 415
     },
     {
       "epoch": 0.7515808491418248,
-      "grad_norm": 0.1448610651899535,
+      "grad_norm": 0.222363917086626,
       "learning_rate": 4.83035384230169e-06,
-      "loss": 0.5537,
+      "loss": 0.5598,
       "step": 416
     },
     {
       "epoch": 0.7533875338753387,
-      "grad_norm": 0.1705391112022685,
+      "grad_norm": 0.23822096081299107,
       "learning_rate": 4.764555358323632e-06,
-      "loss": 0.5756,
+      "loss": 0.5815,
       "step": 417
     },
     {
       "epoch": 0.7551942186088527,
-      "grad_norm": 0.17264531117869172,
+      "grad_norm": 0.2262092084376523,
       "learning_rate": 4.699123384959783e-06,
-      "loss": 0.5576,
+      "loss": 0.5629,
       "step": 418
     },
     {
       "epoch": 0.7570009033423668,
-      "grad_norm": 0.15278513072031447,
+      "grad_norm": 0.23223952114679952,
       "learning_rate": 4.6340602651970304e-06,
-      "loss": 0.559,
+      "loss": 0.5644,
       "step": 419
     },
     {
       "epoch": 0.7588075880758808,
-      "grad_norm": 0.15921165848077162,
+      "grad_norm": 0.23091124005107544,
       "learning_rate": 4.5693683288143884e-06,
-      "loss": 0.5727,
+      "loss": 0.5794,
       "step": 420
     },
     {
       "epoch": 0.7606142728093948,
-      "grad_norm": 0.14827263328330628,
+      "grad_norm": 0.22340107077805454,
       "learning_rate": 4.505049892299517e-06,
-      "loss": 0.5544,
+      "loss": 0.5611,
       "step": 421
     },
     {
       "epoch": 0.7624209575429087,
-      "grad_norm": 0.13887513696660272,
+      "grad_norm": 0.20827372625102367,
       "learning_rate": 4.441107258765827e-06,
-      "loss": 0.5578,
+      "loss": 0.5643,
       "step": 422
     },
     {
       "epoch": 0.7642276422764228,
-      "grad_norm": 0.15863025954427973,
+      "grad_norm": 0.8594750186394492,
       "learning_rate": 4.377542717869959e-06,
-      "loss": 0.5669,
+      "loss": 0.5727,
       "step": 423
     },
     {
       "epoch": 0.7660343270099368,
-      "grad_norm": 0.16888324911277486,
+      "grad_norm": 0.22572761637619165,
       "learning_rate": 4.31435854572983e-06,
-      "loss": 0.5701,
+      "loss": 0.576,
       "step": 424
     },
     {
       "epoch": 0.7678410117434508,
-      "grad_norm": 0.1543410763323266,
+      "grad_norm": 0.22288386378508676,
       "learning_rate": 4.251557004843116e-06,
-      "loss": 0.5489,
+      "loss": 0.5545,
       "step": 425
     },
     {
       "epoch": 0.7696476964769647,
-      "grad_norm": 0.1535360923186684,
+      "grad_norm": 0.22039211202209727,
       "learning_rate": 4.1891403440062495e-06,
-      "loss": 0.5696,
+      "loss": 0.576,
       "step": 426
     },
     {
       "epoch": 0.7714543812104788,
-      "grad_norm": 0.1436847387652519,
+      "grad_norm": 0.24740997697332476,
       "learning_rate": 4.1271107982338745e-06,
-      "loss": 0.541,
+      "loss": 0.5463,
       "step": 427
     },
     {
       "epoch": 0.7732610659439928,
-      "grad_norm": 0.14133645853071267,
+      "grad_norm": 0.21571373017824447,
       "learning_rate": 4.06547058867883e-06,
-      "loss": 0.5616,
+      "loss": 0.5673,
       "step": 428
     },
     {
       "epoch": 0.7750677506775068,
-      "grad_norm": 0.14833855735489168,
+      "grad_norm": 0.21461808803522264,
       "learning_rate": 4.004221922552608e-06,
-      "loss": 0.5567,
+      "loss": 0.5626,
       "step": 429
     },
     {
       "epoch": 0.7768744354110207,
-      "grad_norm": 0.13504301863448828,
+      "grad_norm": 0.21534749337458045,
       "learning_rate": 3.943366993046323e-06,
-      "loss": 0.5573,
+      "loss": 0.5632,
       "step": 430
     },
     {
       "epoch": 0.7786811201445348,
-      "grad_norm": 0.15479562167447294,
+      "grad_norm": 0.22705388420183456,
       "learning_rate": 3.88290797925218e-06,
-      "loss": 0.5665,
+      "loss": 0.5722,
       "step": 431
     },
     {
       "epoch": 0.7804878048780488,
-      "grad_norm": 0.14971545791291554,
+      "grad_norm": 0.20634048423367268,
       "learning_rate": 3.82284704608543e-06,
-      "loss": 0.5618,
+      "loss": 0.5677,
       "step": 432
     },
     {
       "epoch": 0.7822944896115628,
-      "grad_norm": 0.15163059932702447,
+      "grad_norm": 0.22924008526331077,
       "learning_rate": 3.763186344206877e-06,
-      "loss": 0.5694,
+      "loss": 0.5754,
       "step": 433
     },
     {
       "epoch": 0.7841011743450768,
-      "grad_norm": 0.14913417270860077,
+      "grad_norm": 0.19705482767836247,
       "learning_rate": 3.7039280099458373e-06,
-      "loss": 0.5678,
+      "loss": 0.573,
       "step": 434
     },
     {
       "epoch": 0.7859078590785907,
-      "grad_norm": 0.13874898886438794,
+      "grad_norm": 0.24321807481345778,
       "learning_rate": 3.645074165223656e-06,
-      "loss": 0.5457,
+      "loss": 0.552,
       "step": 435
     },
     {
       "epoch": 0.7877145438121048,
-      "grad_norm": 0.15341138739380675,
+      "grad_norm": 0.2557673843285419,
       "learning_rate": 3.586626917477724e-06,
-      "loss": 0.5599,
+      "loss": 0.5663,
       "step": 436
     },
     {
       "epoch": 0.7895212285456188,
-      "grad_norm": 0.1396125925965087,
+      "grad_norm": 0.20741635557070642,
       "learning_rate": 3.528588359586024e-06,
-      "loss": 0.5586,
+      "loss": 0.5646,
       "step": 437
     },
     {
       "epoch": 0.7913279132791328,
-      "grad_norm": 0.1724528227666086,
+      "grad_norm": 0.2688295910485162,
       "learning_rate": 3.4709605697921735e-06,
-      "loss": 0.5499,
+      "loss": 0.5556,
       "step": 438
     },
     {
       "epoch": 0.7931345980126467,
-      "grad_norm": 0.15954102946516008,
+      "grad_norm": 0.46975916458279,
       "learning_rate": 3.413745611631009e-06,
-      "loss": 0.5672,
+      "loss": 0.5736,
       "step": 439
     },
     {
       "epoch": 0.7949412827461608,
-      "grad_norm": 0.13471275186125906,
+      "grad_norm": 0.20860411045588065,
       "learning_rate": 3.356945533854708e-06,
-      "loss": 0.5582,
+      "loss": 0.565,
       "step": 440
     },
     {
       "epoch": 0.7967479674796748,
-      "grad_norm": 0.1277139415264683,
+      "grad_norm": 0.20156415753013815,
       "learning_rate": 3.3005623703594082e-06,
-      "loss": 0.544,
+      "loss": 0.5498,
       "step": 441
     },
     {
       "epoch": 0.7985546522131888,
-      "grad_norm": 0.13269007698962448,
+      "grad_norm": 0.25558258979943804,
       "learning_rate": 3.244598140112404e-06,
-      "loss": 0.5531,
+      "loss": 0.5598,
       "step": 442
     },
     {
       "epoch": 0.8003613369467028,
-      "grad_norm": 0.16876313389901235,
+      "grad_norm": 0.2223887090687036,
       "learning_rate": 3.1890548470798196e-06,
-      "loss": 0.5603,
+      "loss": 0.5666,
       "step": 443
     },
     {
       "epoch": 0.8021680216802168,
-      "grad_norm": 0.14703779234405664,
+      "grad_norm": 0.2805594085283592,
       "learning_rate": 3.1339344801548854e-06,
-      "loss": 0.5662,
+      "loss": 0.5722,
       "step": 444
     },
     {
       "epoch": 0.8039747064137308,
-      "grad_norm": 0.16211778059337975,
+      "grad_norm": 0.22063158100276353,
       "learning_rate": 3.079239013086673e-06,
-      "loss": 0.5541,
+      "loss": 0.5598,
       "step": 445
     },
     {
       "epoch": 0.8057813911472448,
-      "grad_norm": 0.13960967725757564,
+      "grad_norm": 0.20964425176708046,
       "learning_rate": 3.02497040440948e-06,
-      "loss": 0.5521,
+      "loss": 0.5582,
       "step": 446
     },
     {
       "epoch": 0.8075880758807588,
-      "grad_norm": 0.14347409440257566,
+      "grad_norm": 0.19879708507272054,
       "learning_rate": 2.9711305973726404e-06,
-      "loss": 0.5568,
+      "loss": 0.5627,
       "step": 447
     },
     {
       "epoch": 0.8093947606142728,
-      "grad_norm": 0.13685149510690353,
+      "grad_norm": 0.21371375999064623,
       "learning_rate": 2.91772151987099e-06,
-      "loss": 0.5595,
+      "loss": 0.566,
       "step": 448
     },
     {
       "epoch": 0.8112014453477868,
-      "grad_norm": 0.14371214973855934,
+      "grad_norm": 0.21479511072866714,
       "learning_rate": 2.86474508437579e-06,
-      "loss": 0.5541,
+      "loss": 0.5602,
       "step": 449
     },
     {
       "epoch": 0.8130081300813008,
-      "grad_norm": 0.13902608999443639,
+      "grad_norm": 0.2507704066850702,
       "learning_rate": 2.812203187866272e-06,
-      "loss": 0.5651,
+      "loss": 0.5702,
       "step": 450
     },
     {
       "epoch": 0.8148148148148148,
-      "grad_norm": 0.18456478885896846,
+      "grad_norm": 0.21755447137965794,
       "learning_rate": 2.7600977117617027e-06,
-      "loss": 0.566,
+      "loss": 0.5715,
       "step": 451
     },
     {
       "epoch": 0.8166214995483289,
-      "grad_norm": 0.13874263584368351,
+      "grad_norm": 0.20443439016373935,
       "learning_rate": 2.7084305218540217e-06,
-      "loss": 0.5626,
+      "loss": 0.5689,
       "step": 452
     },
     {
       "epoch": 0.8184281842818428,
-      "grad_norm": 0.13915434189224746,
+      "grad_norm": 0.22019364090464516,
       "learning_rate": 2.6572034682410144e-06,
-      "loss": 0.5547,
+      "loss": 0.5601,
       "step": 453
     },
     {
       "epoch": 0.8202348690153568,
-      "grad_norm": 0.14072111536288506,
+      "grad_norm": 0.23616305038551402,
       "learning_rate": 2.60641838526008e-06,
-      "loss": 0.5396,
+      "loss": 0.5441,
       "step": 454
     },
     {
       "epoch": 0.8220415537488708,
-      "grad_norm": 0.13282415327633357,
+      "grad_norm": 0.22116587713371783,
       "learning_rate": 2.5560770914225356e-06,
-      "loss": 0.5625,
+      "loss": 0.5678,
       "step": 455
     },
     {
       "epoch": 0.8238482384823849,
-      "grad_norm": 0.1472812938618023,
+      "grad_norm": 0.2246824985093436,
       "learning_rate": 2.5061813893485087e-06,
-      "loss": 0.5549,
+      "loss": 0.5605,
       "step": 456
     },
     {
       "epoch": 0.8256549232158988,
-      "grad_norm": 0.14223186780386796,
+      "grad_norm": 0.21936692578708686,
       "learning_rate": 2.45673306570239e-06,
-      "loss": 0.5692,
+      "loss": 0.5749,
       "step": 457
     },
     {
       "epoch": 0.8274616079494128,
-      "grad_norm": 0.13382708637823343,
+      "grad_norm": 0.2022912181556414,
       "learning_rate": 2.407733891128841e-06,
-      "loss": 0.5612,
+      "loss": 0.5668,
       "step": 458
     },
     {
       "epoch": 0.8292682926829268,
-      "grad_norm": 0.13648137063446,
+      "grad_norm": 0.20494267015114925,
       "learning_rate": 2.3591856201894124e-06,
-      "loss": 0.5511,
+      "loss": 0.5566,
       "step": 459
     },
     {
       "epoch": 0.8310749774164409,
-      "grad_norm": 0.1547065089258408,
+      "grad_norm": 0.27106860594647675,
       "learning_rate": 2.3110899912996993e-06,
-      "loss": 0.5602,
+      "loss": 0.5663,
       "step": 460
     },
     {
       "epoch": 0.8328816621499548,
-      "grad_norm": 0.14643671146676157,
+      "grad_norm": 0.2046891500263313,
       "learning_rate": 2.263448726667098e-06,
-      "loss": 0.548,
+      "loss": 0.5542,
       "step": 461
     },
     {
       "epoch": 0.8346883468834688,
-      "grad_norm": 0.13651107636976717,
+      "grad_norm": 0.23171892826309606,
       "learning_rate": 2.2162635322291365e-06,
-      "loss": 0.556,
+      "loss": 0.5623,
       "step": 462
     },
     {
       "epoch": 0.8364950316169828,
-      "grad_norm": 0.1368276434381763,
+      "grad_norm": 0.19018606574785138,
       "learning_rate": 2.169536097592401e-06,
-      "loss": 0.5603,
+      "loss": 0.5668,
       "step": 463
     },
     {
       "epoch": 0.8383017163504969,
-      "grad_norm": 0.1445464211298085,
+      "grad_norm": 0.2329448518163566,
       "learning_rate": 2.1232680959720085e-06,
-      "loss": 0.5628,
+      "loss": 0.5692,
       "step": 464
     },
     {
       "epoch": 0.8401084010840109,
-      "grad_norm": 0.1581727260316551,
+      "grad_norm": 0.2221258218802644,
       "learning_rate": 2.0774611841317164e-06,
-      "loss": 0.5728,
+      "loss": 0.5789,
       "step": 465
     },
     {
       "epoch": 0.8419150858175248,
-      "grad_norm": 0.13832300858835891,
+      "grad_norm": 0.2074746474319042,
       "learning_rate": 2.0321170023245835e-06,
-      "loss": 0.5554,
+      "loss": 0.5612,
       "step": 466
     },
     {
       "epoch": 0.8437217705510388,
-      "grad_norm": 0.1326854914012355,
+      "grad_norm": 0.24112788222205025,
       "learning_rate": 1.9872371742342366e-06,
-      "loss": 0.5525,
+      "loss": 0.5588,
       "step": 467
     },
     {
       "epoch": 0.8455284552845529,
-      "grad_norm": 0.1324315694543524,
+      "grad_norm": 0.19795185527219575,
       "learning_rate": 1.942823306916746e-06,
-      "loss": 0.5629,
+      "loss": 0.5693,
       "step": 468
     },
     {
       "epoch": 0.8473351400180669,
-      "grad_norm": 0.14327810150506223,
+      "grad_norm": 0.1977693992081305,
       "learning_rate": 1.8988769907430553e-06,
-      "loss": 0.5522,
+      "loss": 0.5578,
       "step": 469
     },
     {
       "epoch": 0.8491418247515808,
-      "grad_norm": 0.1440779207735543,
+      "grad_norm": 0.2013451884270062,
       "learning_rate": 1.8553997993420495e-06,
-      "loss": 0.5485,
+      "loss": 0.5545,
       "step": 470
     },
     {
       "epoch": 0.8509485094850948,
-      "grad_norm": 0.1582379305959562,
+      "grad_norm": 0.21929485106443544,
       "learning_rate": 1.8123932895441953e-06,
-      "loss": 0.5777,
+      "loss": 0.5831,
       "step": 471
     },
     {
       "epoch": 0.8527551942186089,
-      "grad_norm": 0.13834720570877437,
+      "grad_norm": 0.2069370487233822,
       "learning_rate": 1.7698590013258187e-06,
-      "loss": 0.5486,
+      "loss": 0.5541,
       "step": 472
     },
     {
       "epoch": 0.8545618789521229,
-      "grad_norm": 0.13843120298455933,
+      "grad_norm": 0.21274761276018167,
       "learning_rate": 1.7277984577539224e-06,
-      "loss": 0.5676,
+      "loss": 0.5735,
       "step": 473
     },
     {
       "epoch": 0.8563685636856369,
-      "grad_norm": 0.1279953110251041,
+      "grad_norm": 0.20824350625431515,
       "learning_rate": 1.6862131649316885e-06,
-      "loss": 0.5629,
+      "loss": 0.569,
       "step": 474
     },
     {
       "epoch": 0.8581752484191508,
-      "grad_norm": 0.1585683209243492,
+      "grad_norm": 0.2118004520239169,
       "learning_rate": 1.6451046119445174e-06,
-      "loss": 0.553,
+      "loss": 0.5585,
       "step": 475
     },
     {
       "epoch": 0.8599819331526649,
-      "grad_norm": 0.1466497632190738,
+      "grad_norm": 0.2166189914121313,
       "learning_rate": 1.6044742708067206e-06,
-      "loss": 0.5592,
+      "loss": 0.5655,
       "step": 476
     },
     {
       "epoch": 0.8617886178861789,
-      "grad_norm": 0.1387550564906503,
+      "grad_norm": 0.20548635358142273,
       "learning_rate": 1.5643235964088065e-06,
-      "loss": 0.5564,
+      "loss": 0.5625,
       "step": 477
     },
     {
       "epoch": 0.8635953026196929,
-      "grad_norm": 0.14950588928435643,
+      "grad_norm": 0.3919523574078014,
       "learning_rate": 1.524654026465394e-06,
-      "loss": 0.5671,
+      "loss": 0.5736,
       "step": 478
     },
     {
       "epoch": 0.8654019873532068,
-      "grad_norm": 0.14015867212642105,
+      "grad_norm": 0.2146057433027723,
       "learning_rate": 1.4854669814637145e-06,
-      "loss": 0.5532,
+      "loss": 0.5587,
       "step": 479
     },
     {
       "epoch": 0.8672086720867209,
-      "grad_norm": 0.1262404297566839,
+      "grad_norm": 0.2617349364439417,
       "learning_rate": 1.4467638646127551e-06,
-      "loss": 0.5513,
+      "loss": 0.5577,
       "step": 480
     },
     {
       "epoch": 0.8690153568202349,
-      "grad_norm": 0.12824154810533317,
+      "grad_norm": 0.22872811738070856,
       "learning_rate": 1.4085460617930196e-06,
-      "loss": 0.5622,
+      "loss": 0.5689,
       "step": 481
     },
     {
       "epoch": 0.8708220415537489,
-      "grad_norm": 0.1345600337122948,
+      "grad_norm": 0.2796054869270635,
       "learning_rate": 1.3708149415068827e-06,
-      "loss": 0.5661,
+      "loss": 0.5725,
       "step": 482
     },
     {
       "epoch": 0.8726287262872628,
-      "grad_norm": 0.13838894371767235,
+      "grad_norm": 0.19160336089610616,
       "learning_rate": 1.3335718548296194e-06,
-      "loss": 0.5504,
+      "loss": 0.5563,
       "step": 483
     },
     {
       "epoch": 0.8744354110207768,
-      "grad_norm": 0.13576491587772027,
+      "grad_norm": 0.2012810744492221,
       "learning_rate": 1.2968181353609854e-06,
-      "loss": 0.5472,
+      "loss": 0.5535,
       "step": 484
     },
     {
       "epoch": 0.8762420957542909,
-      "grad_norm": 0.14339241553346477,
+      "grad_norm": 0.21604349922209895,
       "learning_rate": 1.2605550991775033e-06,
-      "loss": 0.5576,
+      "loss": 0.5629,
       "step": 485
     },
     {
       "epoch": 0.8780487804878049,
-      "grad_norm": 0.131210165981635,
+      "grad_norm": 0.20846295082414001,
       "learning_rate": 1.2247840447852993e-06,
-      "loss": 0.5633,
+      "loss": 0.5702,
       "step": 486
     },
     {
       "epoch": 0.8798554652213189,
-      "grad_norm": 0.12636963673731474,
+      "grad_norm": 0.24538983066246176,
       "learning_rate": 1.1895062530736412e-06,
-      "loss": 0.5515,
+      "loss": 0.5577,
       "step": 487
     },
     {
       "epoch": 0.8816621499548328,
-      "grad_norm": 0.12776691333292706,
+      "grad_norm": 0.2027134974758516,
       "learning_rate": 1.1547229872690407e-06,
-      "loss": 0.5505,
+      "loss": 0.5567,
       "step": 488
     },
     {
       "epoch": 0.8834688346883469,
-      "grad_norm": 0.12702102055757536,
+      "grad_norm": 0.18716638337399608,
       "learning_rate": 1.1204354928900496e-06,
-      "loss": 0.5602,
+      "loss": 0.5663,
       "step": 489
     },
     {
       "epoch": 0.8852755194218609,
-      "grad_norm": 0.13494740824410642,
+      "grad_norm": 0.21298790397676395,
       "learning_rate": 1.086644997702635e-06,
-      "loss": 0.5571,
+      "loss": 0.5625,
       "step": 490
     },
     {
       "epoch": 0.8870822041553749,
-      "grad_norm": 0.12513086965807382,
+      "grad_norm": 0.18568745233434789,
       "learning_rate": 1.0533527116762298e-06,
-      "loss": 0.5483,
+      "loss": 0.5551,
       "step": 491
     },
     {
       "epoch": 0.8888888888888888,
-      "grad_norm": 0.1320750230179298,
+      "grad_norm": 0.20538517713781662,
       "learning_rate": 1.0205598269404026e-06,
-      "loss": 0.55,
+      "loss": 0.5556,
       "step": 492
     },
     {
       "epoch": 0.8906955736224029,
-      "grad_norm": 0.12753125629115394,
+      "grad_norm": 0.21478395379619078,
       "learning_rate": 9.882675177421662e-07,
-      "loss": 0.5444,
+      "loss": 0.5506,
       "step": 493
     },
     {
       "epoch": 0.8925022583559169,
-      "grad_norm": 0.12278292224063106,
+      "grad_norm": 0.20451931843479138,
       "learning_rate": 9.56476940403942e-07,
-      "loss": 0.5617,
+      "loss": 0.5672,
       "step": 494
     },
     {
       "epoch": 0.8943089430894309,
-      "grad_norm": 0.13711394488670506,
+      "grad_norm": 0.21150501837407576,
       "learning_rate": 9.251892332821388e-07,
-      "loss": 0.5584,
+      "loss": 0.5641,
       "step": 495
     },
     {
       "epoch": 0.8961156278229448,
-      "grad_norm": 0.13665462597200176,
+      "grad_norm": 0.20494840081387994,
       "learning_rate": 8.944055167263982e-07,
-      "loss": 0.5706,
+      "loss": 0.5768,
       "step": 496
     },
     {
       "epoch": 0.8979223125564589,
-      "grad_norm": 0.14101341052067656,
+      "grad_norm": 0.23208540837310512,
       "learning_rate": 8.641268930394791e-07,
-      "loss": 0.5578,
+      "loss": 0.5637,
       "step": 497
     },
     {
       "epoch": 0.8997289972899729,
-      "grad_norm": 0.1371088435731868,
+      "grad_norm": 0.43670103718238795,
       "learning_rate": 8.34354446437785e-07,
-      "loss": 0.5571,
+      "loss": 0.5631,
       "step": 498
     },
     {
       "epoch": 0.9015356820234869,
-      "grad_norm": 0.15044968278360257,
+      "grad_norm": 0.22601413932793338,
       "learning_rate": 8.050892430125362e-07,
-      "loss": 0.5662,
+      "loss": 0.5728,
       "step": 499
     },
     {
       "epoch": 0.9033423667570009,
-      "grad_norm": 0.13190920385566557,
+      "grad_norm": 0.19668433127008136,
       "learning_rate": 7.763323306916037e-07,
-      "loss": 0.565,
+      "loss": 0.5715,
       "step": 500
     },
     {
       "epoch": 0.9051490514905149,
-      "grad_norm": 0.1330272682028843,
+      "grad_norm": 0.2755002021494048,
       "learning_rate": 7.48084739201978e-07,
-      "loss": 0.5629,
+      "loss": 0.5691,
       "step": 501
     },
     {
       "epoch": 0.9069557362240289,
-      "grad_norm": 0.12596324805796347,
+      "grad_norm": 0.3171350403732765,
       "learning_rate": 7.20347480032893e-07,
-      "loss": 0.5575,
+      "loss": 0.5636,
       "step": 502
     },
     {
       "epoch": 0.9087624209575429,
-      "grad_norm": 0.12269256945275395,
+      "grad_norm": 0.19229422822362555,
       "learning_rate": 6.931215463996204e-07,
-      "loss": 0.5536,
+      "loss": 0.5592,
       "step": 503
     },
     {
       "epoch": 0.9105691056910569,
-      "grad_norm": 0.12397511410433176,
+      "grad_norm": 0.27683701338372546,
       "learning_rate": 6.664079132078881e-07,
-      "loss": 0.5693,
+      "loss": 0.5751,
       "step": 504
     },
     {
       "epoch": 0.912375790424571,
-      "grad_norm": 0.12830241052252903,
+      "grad_norm": 0.2809384935828692,
       "learning_rate": 6.402075370189914e-07,
-      "loss": 0.541,
+      "loss": 0.5459,
       "step": 505
     },
     {
       "epoch": 0.9141824751580849,
-      "grad_norm": 0.13190860459526532,
+      "grad_norm": 0.19763837414051177,
       "learning_rate": 6.145213560155167e-07,
-      "loss": 0.5613,
+      "loss": 0.5665,
       "step": 506
     },
     {
       "epoch": 0.9159891598915989,
-      "grad_norm": 0.13145720432300395,
+      "grad_norm": 0.19330852387204572,
       "learning_rate": 5.893502899677622e-07,
-      "loss": 0.5597,
+      "loss": 0.5655,
       "step": 507
     },
     {
       "epoch": 0.9177958446251129,
-      "grad_norm": 0.1390726475782522,
+      "grad_norm": 0.18671045701948052,
       "learning_rate": 5.646952402007993e-07,
-      "loss": 0.5523,
+      "loss": 0.5578,
       "step": 508
     },
     {
       "epoch": 0.919602529358627,
-      "grad_norm": 0.13621784896319666,
+      "grad_norm": 0.1935404379381202,
       "learning_rate": 5.405570895622014e-07,
-      "loss": 0.5716,
+      "loss": 0.5771,
       "step": 509
     },
     {
       "epoch": 0.9214092140921409,
-      "grad_norm": 0.1303455229261674,
+      "grad_norm": 0.2266010335193308,
       "learning_rate": 5.169367023904226e-07,
-      "loss": 0.5547,
+      "loss": 0.5612,
       "step": 510
     },
     {
       "epoch": 0.9232158988256549,
-      "grad_norm": 0.12456382952524406,
+      "grad_norm": 0.18970071959766807,
       "learning_rate": 4.938349244838541e-07,
-      "loss": 0.5452,
+      "loss": 0.5503,
       "step": 511
     },
     {
       "epoch": 0.9250225835591689,
-      "grad_norm": 0.13033855895306248,
+      "grad_norm": 0.18751450133804787,
       "learning_rate": 4.7125258307053385e-07,
-      "loss": 0.5725,
+      "loss": 0.5788,
       "step": 512
     },
     {
       "epoch": 0.926829268292683,
-      "grad_norm": 0.12446562507353466,
+      "grad_norm": 0.20331269139776248,
       "learning_rate": 4.491904867785329e-07,
-      "loss": 0.5548,
+      "loss": 0.5604,
       "step": 513
     },
     {
       "epoch": 0.928635953026197,
-      "grad_norm": 0.13470648707570285,
+      "grad_norm": 0.23016504488814535,
       "learning_rate": 4.276494256069874e-07,
-      "loss": 0.5512,
+      "loss": 0.5567,
       "step": 514
     },
     {
       "epoch": 0.9304426377597109,
-      "grad_norm": 0.12820525405710728,
+      "grad_norm": 0.19554882132706672,
       "learning_rate": 4.0663017089782593e-07,
-      "loss": 0.5567,
+      "loss": 0.5625,
       "step": 515
     },
     {
       "epoch": 0.9322493224932249,
-      "grad_norm": 0.1285059094961294,
+      "grad_norm": 0.2141915434693991,
       "learning_rate": 3.8613347530813527e-07,
-      "loss": 0.5649,
+      "loss": 0.5705,
       "step": 516
     },
     {
       "epoch": 0.934056007226739,
-      "grad_norm": 0.12804520061809396,
+      "grad_norm": 0.19325511174233914,
       "learning_rate": 3.6616007278320985e-07,
-      "loss": 0.5529,
+      "loss": 0.5589,
       "step": 517
     },
     {
       "epoch": 0.935862691960253,
-      "grad_norm": 0.13304365719456168,
+      "grad_norm": 0.2867842616135943,
       "learning_rate": 3.467106785302865e-07,
-      "loss": 0.5545,
+      "loss": 0.5597,
       "step": 518
     },
     {
       "epoch": 0.9376693766937669,
-      "grad_norm": 0.1361684924628673,
+      "grad_norm": 0.1965457021012975,
       "learning_rate": 3.277859889929147e-07,
-      "loss": 0.5444,
+      "loss": 0.5506,
       "step": 519
     },
     {
       "epoch": 0.9394760614272809,
-      "grad_norm": 0.1308177622907266,
+      "grad_norm": 0.1983487205786454,
       "learning_rate": 3.0938668182603656e-07,
-      "loss": 0.5538,
+      "loss": 0.5599,
       "step": 520
     },
     {
       "epoch": 0.941282746160795,
-      "grad_norm": 0.12589685134140113,
+      "grad_norm": 0.2016483308305746,
       "learning_rate": 2.9151341587169976e-07,
-      "loss": 0.5614,
+      "loss": 0.5676,
       "step": 521
     },
     {
       "epoch": 0.943089430894309,
-      "grad_norm": 0.1309203120133497,
+      "grad_norm": 0.18584462016114534,
       "learning_rate": 2.7416683113548945e-07,
-      "loss": 0.5568,
+      "loss": 0.5628,
       "step": 522
     },
     {
       "epoch": 0.9448961156278229,
-      "grad_norm": 0.11901547198938607,
+      "grad_norm": 0.1978444406310148,
       "learning_rate": 2.5734754876358846e-07,
-      "loss": 0.5478,
+      "loss": 0.5537,
       "step": 523
     },
     {
       "epoch": 0.9467028003613369,
-      "grad_norm": 0.13909322877084576,
+      "grad_norm": 0.36806346840727133,
       "learning_rate": 2.41056171020555e-07,
-      "loss": 0.5577,
+      "loss": 0.5638,
       "step": 524
     },
     {
       "epoch": 0.948509485094851,
-      "grad_norm": 0.1250829734778908,
+      "grad_norm": 0.21280062488260573,
       "learning_rate": 2.2529328126774496e-07,
-      "loss": 0.5511,
+      "loss": 0.557,
       "step": 525
     },
     {
       "epoch": 0.950316169828365,
-      "grad_norm": 0.12347363002885929,
+      "grad_norm": 0.19466678941893495,
       "learning_rate": 2.1005944394242694e-07,
-      "loss": 0.5555,
+      "loss": 0.5604,
       "step": 526
     },
     {
       "epoch": 0.952122854561879,
-      "grad_norm": 0.14522071992485233,
+      "grad_norm": 0.21086407350506756,
       "learning_rate": 1.9535520453756506e-07,
-      "loss": 0.5556,
+      "loss": 0.562,
       "step": 527
     },
     {
       "epoch": 0.9539295392953929,
-      "grad_norm": 0.12689680120062793,
+      "grad_norm": 0.30098595384873267,
       "learning_rate": 1.8118108958229795e-07,
-      "loss": 0.5536,
+      "loss": 0.5596,
       "step": 528
     },
     {
       "epoch": 0.955736224028907,
-      "grad_norm": 0.12751446868474398,
+      "grad_norm": 0.19293881839582736,
       "learning_rate": 1.6753760662307217e-07,
-      "loss": 0.56,
+      "loss": 0.566,
       "step": 529
     },
     {
       "epoch": 0.957542908762421,
-      "grad_norm": 0.1377057005084481,
+      "grad_norm": 0.2178986422135447,
       "learning_rate": 1.5442524420547666e-07,
-      "loss": 0.5515,
+      "loss": 0.5576,
       "step": 530
     },
     {
       "epoch": 0.959349593495935,
-      "grad_norm": 0.12467052220580256,
+      "grad_norm": 0.2742424638712097,
       "learning_rate": 1.4184447185674354e-07,
-      "loss": 0.5642,
+      "loss": 0.5712,
       "step": 531
     },
     {
       "epoch": 0.9611562782294489,
-      "grad_norm": 0.12474201474426573,
+      "grad_norm": 0.20441449476859883,
       "learning_rate": 1.297957400689398e-07,
-      "loss": 0.5532,
+      "loss": 0.5589,
       "step": 532
     },
     {
       "epoch": 0.9629629629629629,
-      "grad_norm": 0.1378124489625079,
+      "grad_norm": 0.2028488844322694,
       "learning_rate": 1.1827948028283353e-07,
-      "loss": 0.5578,
+      "loss": 0.564,
       "step": 533
     },
     {
       "epoch": 0.964769647696477,
-      "grad_norm": 0.12241633221352105,
+      "grad_norm": 0.19641796792115376,
       "learning_rate": 1.0729610487243968e-07,
-      "loss": 0.554,
+      "loss": 0.56,
       "step": 534
     },
     {
       "epoch": 0.966576332429991,
-      "grad_norm": 0.12161233750054955,
+      "grad_norm": 0.1961644503643991,
       "learning_rate": 9.684600713026847e-08,
-      "loss": 0.5549,
+      "loss": 0.5613,
       "step": 535
     },
     {
       "epoch": 0.9683830171635049,
-      "grad_norm": 0.12789778109710662,
+      "grad_norm": 0.19348594352650886,
       "learning_rate": 8.692956125322837e-08,
-      "loss": 0.5637,
+      "loss": 0.5693,
       "step": 536
     },
     {
       "epoch": 0.9701897018970189,
-      "grad_norm": 0.12936381253059515,
+      "grad_norm": 0.26854880925929886,
       "learning_rate": 7.754712232922678e-08,
-      "loss": 0.5699,
+      "loss": 0.5761,
       "step": 537
     },
     {
       "epoch": 0.971996386630533,
-      "grad_norm": 0.14000238842546384,
+      "grad_norm": 0.20691432526444847,
       "learning_rate": 6.869902632447022e-08,
-      "loss": 0.5537,
+      "loss": 0.5591,
       "step": 538
     },
     {
       "epoch": 0.973803071364047,
-      "grad_norm": 0.12479645822662797,
+      "grad_norm": 0.18552138088036305,
       "learning_rate": 6.038559007141397e-08,
-      "loss": 0.5509,
+      "loss": 0.5564,
       "step": 539
     },
     {
       "epoch": 0.975609756097561,
-      "grad_norm": 0.12949980004311576,
+      "grad_norm": 0.19379985533649824,
       "learning_rate": 5.260711125743445e-08,
-      "loss": 0.5512,
+      "loss": 0.5575,
       "step": 540
     },
     {
       "epoch": 0.9774164408310749,
-      "grad_norm": 0.12362590360829781,
+      "grad_norm": 0.1915865386485388,
       "learning_rate": 4.536386841415774e-08,
-      "loss": 0.56,
+      "loss": 0.5667,
       "step": 541
     },
     {
       "epoch": 0.979223125564589,
-      "grad_norm": 0.1318525107421328,
+      "grad_norm": 0.2671463373041046,
       "learning_rate": 3.8656120907489266e-08,
-      "loss": 0.5524,
+      "loss": 0.5582,
       "step": 542
     },
     {
       "epoch": 0.981029810298103,
-      "grad_norm": 0.12341218300550963,
+      "grad_norm": 0.1885606186425532,
       "learning_rate": 3.24841089283312e-08,
-      "loss": 0.5661,
+      "loss": 0.5718,
       "step": 543
     },
     {
       "epoch": 0.982836495031617,
-      "grad_norm": 0.1313814849485639,
+      "grad_norm": 0.2328688039271822,
       "learning_rate": 2.684805348397268e-08,
-      "loss": 0.5653,
+      "loss": 0.5714,
       "step": 544
     },
     {
       "epoch": 0.9846431797651309,
-      "grad_norm": 0.1246027162262715,
+      "grad_norm": 0.1940328082205833,
       "learning_rate": 2.1748156390186125e-08,
-      "loss": 0.5679,
+      "loss": 0.5744,
       "step": 545
     },
     {
       "epoch": 0.986449864498645,
-      "grad_norm": 0.12930758713143098,
+      "grad_norm": 0.20170122816300584,
       "learning_rate": 1.718460026399138e-08,
-      "loss": 0.5575,
+      "loss": 0.5632,
       "step": 546
     },
     {
       "epoch": 0.988256549232159,
-      "grad_norm": 0.12478369737661307,
+      "grad_norm": 0.19578981489849026,
       "learning_rate": 1.3157548517124251e-08,
-      "loss": 0.5697,
+      "loss": 0.5759,
       "step": 547
     },
     {
       "epoch": 0.990063233965673,
-      "grad_norm": 0.1311592043703811,
+      "grad_norm": 0.24051446818881467,
       "learning_rate": 9.66714535017954e-09,
-      "loss": 0.5666,
+      "loss": 0.5724,
       "step": 548
     },
     {
       "epoch": 0.991869918699187,
-      "grad_norm": 0.12850976653433435,
+      "grad_norm": 0.18847855702618435,
       "learning_rate": 6.7135157474501604e-09,
-      "loss": 0.5611,
+      "loss": 0.5674,
       "step": 549
     },
     {
       "epoch": 0.993676603432701,
-      "grad_norm": 0.1316930656860681,
+      "grad_norm": 0.2042654954909486,
       "learning_rate": 4.29676547244906e-09,
-      "loss": 0.5511,
+      "loss": 0.5569,
       "step": 550
     },
     {
       "epoch": 0.995483288166215,
-      "grad_norm": 0.12869152675188197,
+      "grad_norm": 0.1884805767867668,
       "learning_rate": 2.416981064127244e-09,
-      "loss": 0.57,
+      "loss": 0.5764,
       "step": 551
     },
     {
       "epoch": 0.997289972899729,
-      "grad_norm": 0.13230005405798148,
+      "grad_norm": 0.25949043814500883,
       "learning_rate": 1.074229833769591e-09,
-      "loss": 0.5523,
+      "loss": 0.5574,
       "step": 552
     },
     {
       "epoch": 0.999096657633243,
-      "grad_norm": 0.13151662587114868,
+      "grad_norm": 0.21033436815555134,
       "learning_rate": 2.6855986258844666e-10,
-      "loss": 0.5559,
+      "loss": 0.5612,
       "step": 553
     }
   ],