Upload weights and configs - Run 20251012_060013

weights/best_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:639dff7e61f814a6d11687f482bfc5f67618578b39c13326a1ee72ea3de462d6
+oid sha256:da3fe18b58dcc60496b6dae7a4b533511f6ac5f4075a4a2701aa879481b2a188
 size 2628344

weights/best_model_metadata.json

@@ -1,46 +1,46 @@
 {
-  "epoch": 8,
+  "epoch": 9,
   "optimizer_state_dict": {
     "state": {
       "0": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([[-6.4025e-05, -5.4984e-05,  1.2083e-05,  ...,  3.5269e-05,\n         -1.0064e-04,  1.1238e-04],\n        [-6.1181e-04,  5.2606e-04,  6.2456e-05,  ..., -2.2401e-04,\n         -2.8836e-04, -5.2963e-05],\n        [-1.3835e-05,  1.2044e-05, -4.2223e-05,  ...,  3.5661e-05,\n          3.7672e-05, -2.7067e-05],\n        ...,\n        [ 2.3016e-04, -7.2059e-05, -4.2696e-05,  ...,  1.1419e-06,\n         -2.8614e-05,  6.3013e-05],\n        [ 9.8212e-05,  1.6402e-05, -9.2308e-06,  ..., -2.3917e-05,\n         -7.1468e-05, -7.9436e-06],\n        [ 8.3501e-05,  2.7023e-05,  4.5033e-05,  ...,  4.4794e-05,\n          8.6830e-05, -1.7564e-05]], device='cuda:0')",
-        "exp_avg_sq": "tensor([[1.8269e-07, 2.8111e-07, 5.0383e-08,  ..., 4.1391e-08, 8.4271e-08,\n         5.3235e-08],\n        [2.6887e-07, 4.8796e-07, 5.2036e-08,  ..., 4.0222e-08, 1.4556e-07,\n         7.2980e-08],\n        [2.2098e-07, 4.8053e-07, 6.4593e-08,  ..., 6.0801e-08, 5.5337e-08,\n         4.5989e-08],\n        ...,\n        [1.4012e-07, 3.9894e-07, 4.6741e-08,  ..., 5.1466e-08, 4.2373e-08,\n         4.1448e-08],\n        [2.7050e-07, 2.0343e-07, 4.8844e-08,  ..., 6.6012e-08, 5.2131e-08,\n         4.2363e-08],\n        [1.2194e-07, 1.4229e-07, 3.5382e-08,  ..., 5.0352e-08, 4.6794e-08,\n         3.8750e-08]], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([[ 1.5781e-04, -3.5266e-04, -9.7583e-05,  ...,  4.1940e-05,\n          5.6975e-05,  5.5798e-05],\n        [ 9.8742e-05, -1.3119e-04,  1.3677e-05,  ..., -1.9366e-05,\n          1.0349e-04,  5.0592e-05],\n        [-2.9130e-04,  1.9895e-04, -9.0260e-05,  ..., -9.9918e-05,\n         -1.8629e-04, -6.0961e-05],\n        ...,\n        [ 5.7664e-05,  4.7397e-05,  1.2529e-04,  ...,  2.8799e-05,\n         -5.4763e-05, -6.6750e-06],\n        [ 6.3787e-05, -2.1633e-05,  6.6595e-05,  ...,  2.7704e-05,\n          8.9652e-05, -4.2091e-06],\n        [ 1.5883e-04, -1.1030e-04,  6.1721e-05,  ...,  1.1175e-04,\n         -8.8536e-06,  8.7643e-05]], device='cuda:0')",
+        "exp_avg_sq": "tensor([[1.6698e-07, 2.8303e-07, 5.0525e-08,  ..., 4.0017e-08, 8.8370e-08,\n         5.3000e-08],\n        [2.5771e-07, 4.5601e-07, 5.9020e-08,  ..., 4.0411e-08, 1.4522e-07,\n         7.0853e-08],\n        [2.3775e-07, 4.7696e-07, 6.2997e-08,  ..., 6.3045e-08, 5.9060e-08,\n         4.6280e-08],\n        ...,\n        [1.4690e-07, 4.0507e-07, 4.7863e-08,  ..., 4.8895e-08, 4.3604e-08,\n         3.9252e-08],\n        [2.7245e-07, 2.1035e-07, 4.9765e-08,  ..., 6.8004e-08, 5.1936e-08,\n         4.4202e-08],\n        [1.1928e-07, 1.4219e-07, 3.6373e-08,  ..., 5.2237e-08, 4.5473e-08,\n         3.9891e-08]], device='cuda:0')"
       },
       "1": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([-6.4987e-04, -1.3535e-02,  1.2573e-03, -6.5791e-04, -5.5231e-03,\n        -3.6399e-04, -5.6482e-04,  3.4371e-03,  3.7432e-03, -3.5744e-03,\n        -2.8670e-03,  4.1944e-05,  4.0621e-03, -1.7573e-03,  3.0126e-04,\n        -1.7162e-03,  5.3238e-03,  5.8619e-04,  5.4343e-03, -1.8391e-03,\n        -4.0455e-04, -1.4090e-03,  9.1251e-04,  8.3566e-04, -4.9397e-03,\n        -1.0591e-03,  1.6060e-03, -1.2161e-03,  2.8320e-04, -6.0574e-03,\n        -1.1296e-03, -8.4397e-03,  8.7225e-04, -1.0327e-03, -1.9377e-03,\n         1.1576e-04,  2.6625e-04,  1.3670e-03, -1.0076e-03, -1.8584e-03,\n         1.5999e-03, -3.7180e-04, -2.9035e-03, -2.9538e-03, -3.3774e-03,\n        -2.1114e-03,  7.1032e-05,  3.7267e-03,  1.0947e-03,  6.4281e-04,\n        -5.1595e-04, -1.4746e-03, -5.3556e-04, -1.4137e-03, -6.6869e-04,\n         4.8831e-04,  4.0058e-03, -3.2143e-03,  2.4246e-03,  4.0907e-03,\n         1.7049e-03,  1.9003e-03, -4.9165e-03, -2.2449e-03,  2.8681e-03,\n         2.5851e-03, -2.3767e-03, -2.6131e-03, -1.0220e-03,  9.9959e-04,\n         2.0318e-03, -1.3424e-03, -5.0590e-03,  4.2861e-04,  1.3476e-03,\n        -7.7529e-04, -2.4318e-03,  5.7784e-04,  1.8398e-03, -3.8252e-05,\n        -4.4108e-03, -2.6385e-03, -2.4669e-03,  3.9434e-04, -1.0643e-03,\n        -6.1989e-04, -4.4558e-03,  3.2091e-03,  1.2780e-03,  2.1135e-03,\n        -6.6324e-03, -3.3958e-03,  6.3886e-04, -2.2225e-03,  6.8824e-03,\n        -1.3394e-03, -2.4285e-03,  2.0074e-03, -7.7845e-04, -4.5406e-03,\n         9.0649e-04, -2.1313e-03, -1.0848e-03, -2.3967e-03, -3.9903e-03,\n         4.0038e-03, -3.0237e-03,  2.8845e-03, -2.3969e-03,  5.2510e-05,\n         6.0553e-04, -1.6673e-03, -2.5653e-04,  2.8108e-04, -5.5730e-04,\n         4.2882e-03, -1.4809e-04,  1.7556e-04,  3.4074e-04, -1.3873e-04,\n         2.5531e-03,  3.5905e-03,  9.0629e-04,  1.1274e-03, -4.5558e-03,\n        -6.7897e-03, -3.1809e-04, -9.3280e-04,  3.8896e-03, -2.2102e-04,\n        -9.3409e-04, -2.7715e-03, -1.4590e-03, -1.3147e-04,  5.1617e-05,\n        -9.5770e-04,  1.4158e-03,  1.3385e-03, -8.0067e-04, -5.5641e-03,\n         2.5734e-03, -2.3095e-03, -2.6607e-03, -3.3851e-03, -2.1216e-03,\n        -2.4541e-03, -1.9666e-04, -1.0192e-03,  2.3320e-05, -3.2565e-03,\n        -4.5256e-03, -1.0836e-03, -1.2236e-03,  1.9268e-03, -1.2424e-03,\n        -4.1846e-03,  1.6557e-03,  1.8472e-03,  5.1666e-03,  2.1120e-03,\n         7.1577e-04, -1.5775e-03, -4.9247e-04,  2.0997e-04, -5.3849e-04,\n         1.3241e-03,  1.0897e-03,  1.2397e-03,  3.2353e-04,  4.7124e-04,\n         1.1061e-03,  1.9082e-03, -1.5052e-03,  6.5254e-04,  7.7406e-06,\n        -4.3283e-04,  4.4846e-03,  1.0456e-03,  4.2589e-03, -4.8122e-03,\n         3.5219e-03,  1.4128e-03, -1.8019e-03,  2.0805e-03,  1.1614e-03,\n        -3.3549e-03,  8.6528e-04, -1.0777e-03, -2.6108e-03,  1.9170e-03,\n         2.0967e-03, -1.4040e-03,  5.4892e-03, -7.0758e-04,  5.2032e-04,\n         2.7227e-03,  2.3475e-03, -9.8022e-04, -3.5715e-04,  2.1807e-03,\n         7.6662e-04, -6.1029e-03,  9.4796e-05, -9.0100e-03, -2.0703e-04,\n         2.4527e-04, -1.5332e-03, -6.1512e-04,  1.5478e-03, -3.8316e-03,\n         4.9976e-04,  1.8426e-03,  3.6613e-03,  1.2577e-03,  3.3443e-03,\n         3.7792e-03,  2.5343e-03,  1.4302e-04,  5.8838e-04,  4.5311e-03,\n         1.5365e-03,  4.8269e-03,  1.1041e-04, -1.7578e-03,  1.7293e-03,\n         3.2841e-03, -1.5400e-04,  1.0667e-03,  2.4802e-03, -1.1687e-03,\n         7.2003e-05,  1.3605e-04,  1.2642e-03,  1.4252e-03, -6.6624e-04,\n         5.4636e-04,  3.0512e-03, -1.1773e-03,  9.3631e-04, -2.3606e-03,\n         5.4453e-03, -2.7824e-03,  2.1358e-03, -8.8854e-04, -1.8944e-03,\n        -2.0677e-03, -6.0333e-04,  3.9044e-03, -3.1116e-04,  9.2369e-04,\n        -3.4924e-03, -1.0965e-03, -2.4552e-03,  2.6810e-04, -2.5463e-04,\n         2.0702e-03,  4.4322e-04, -5.9029e-03, -2.0743e-03,  1.3176e-03,\n        -2.4826e-05, -3.4424e-04,  2.6223e-03,  9.4406e-04,  8.7087e-04,\n         1.4140e-03,  1.3032e-03, -3.5142e-03, -6.8172e-05, -1.6200e-03,\n        -2.5520e-03,  2.4010e-03,  1.4463e-03,  1.3729e-04,  1.7187e-03,\n         3.5223e-03,  2.7355e-03, -9.8063e-04, -3.7862e-03, -7.4730e-04,\n         1.4966e-03,  2.4800e-03,  2.0105e-03,  2.7326e-04,  4.9518e-03,\n         5.1635e-03,  1.4258e-03, -3.8533e-05,  2.0503e-03,  4.8432e-04,\n         2.7832e-03,  1.6141e-03, -9.7171e-06, -6.4220e-04, -7.2672e-03,\n        -2.1605e-03, -3.8668e-03,  1.9258e-03,  3.8126e-03,  4.5958e-03,\n        -1.8867e-04,  2.4431e-03,  2.5902e-03, -7.4286e-03,  1.4470e-03,\n        -3.1754e-03, -6.1963e-03, -4.4241e-05, -1.2219e-03, -5.1569e-04,\n         5.6235e-04,  7.3495e-05, -3.9182e-04, -2.4586e-03,  2.2537e-03,\n         5.2585e-04, -6.7306e-04,  1.1516e-03, -6.3138e-03, -6.1228e-04,\n        -5.6314e-04,  2.0937e-03,  2.3887e-05, -3.9378e-03,  4.5113e-04,\n         6.9726e-04, -3.1392e-03, -4.1529e-03,  4.1541e-05,  1.9527e-03,\n         2.4701e-03,  4.6276e-03, -4.6993e-04, -1.4593e-03,  3.6660e-03,\n        -3.0804e-03, -2.7708e-03, -3.6792e-04, -3.7439e-03, -4.3573e-04,\n        -9.3716e-04, -4.4792e-04,  2.7432e-03,  1.2436e-03,  7.0712e-04,\n        -6.5514e-03,  2.2349e-03, -1.7639e-03, -6.7714e-04,  2.6556e-03,\n         2.6999e-03, -1.3574e-03,  1.4930e-03, -3.6532e-03, -4.4568e-03,\n        -1.7960e-03, -2.2475e-03, -4.9731e-04, -2.3124e-03,  1.2356e-03,\n         6.3104e-03,  1.4537e-03, -9.7431e-04, -2.4065e-03,  1.3232e-04,\n        -1.2747e-03, -1.9836e-03,  9.3327e-03, -5.0278e-03,  5.6487e-04,\n        -4.3260e-03,  1.7001e-03, -4.6955e-04, -2.6995e-04, -3.2107e-03,\n        -5.1682e-04,  3.6555e-04,  3.5020e-03,  4.2030e-03,  3.1604e-03,\n        -7.2903e-04, -2.2251e-03,  1.4401e-03,  2.9910e-03,  7.3149e-04,\n        -1.4167e-03,  2.9373e-04, -2.7491e-03,  1.3375e-03, -7.8760e-03,\n         1.8148e-03,  2.1893e-03,  4.1376e-03,  1.2380e-03, -1.9388e-03,\n        -4.1796e-03, -3.3875e-04, -2.2825e-03,  1.5630e-03, -7.4003e-03,\n        -2.1383e-03,  2.8884e-03,  2.1469e-03, -8.5662e-04, -1.5254e-03,\n         2.4844e-03, -6.5807e-05,  7.5317e-04,  2.9965e-03,  5.5822e-04,\n         1.9286e-03, -5.0742e-05,  3.4163e-03,  2.2084e-03,  5.7377e-04,\n        -4.3959e-04,  2.8400e-03,  1.1998e-03,  6.1082e-03,  5.1326e-03,\n         9.2500e-04,  5.0460e-03, -1.4308e-03, -2.6501e-03,  1.6726e-03,\n         1.7697e-03, -1.6872e-05,  1.0886e-03,  4.6791e-03,  2.5985e-04,\n        -4.6035e-03,  7.3891e-04, -1.9553e-03,  2.7678e-03,  2.9167e-03,\n         4.0055e-05, -1.0356e-02, -5.4378e-03,  8.8839e-04, -1.3491e-03,\n         7.6902e-04, -9.1812e-04,  2.9943e-03,  1.4666e-03, -4.6294e-04,\n        -8.7309e-04,  2.2775e-04,  6.4789e-04, -7.2850e-05,  1.7497e-03,\n         2.0643e-03, -7.8843e-04, -2.8440e-03,  1.7595e-03,  2.4526e-03,\n        -2.5940e-03,  2.5147e-03,  1.9497e-03, -2.3792e-03,  2.0838e-03,\n        -1.0004e-03, -4.6084e-04, -1.5070e-03, -3.1268e-03, -9.0118e-05,\n         1.2573e-03,  8.1070e-03,  3.2379e-03,  1.0477e-03,  8.8266e-03,\n        -4.9239e-03, -2.4421e-03, -1.5647e-03, -8.0027e-04,  2.1143e-03,\n         1.2663e-03,  1.1784e-03,  9.0763e-04, -4.7574e-04,  3.0364e-04,\n         3.0000e-03,  1.2367e-02,  2.7381e-03, -1.5734e-03,  3.8696e-03,\n        -3.3985e-03, -4.3234e-03, -2.2549e-03,  1.7785e-03,  1.6328e-03,\n         1.5920e-03,  9.6109e-04,  1.7402e-03,  3.0145e-03,  2.6775e-03,\n         1.4967e-03,  1.5786e-03,  5.9222e-04,  1.9967e-03,  1.6153e-04,\n        -2.7403e-03, -1.9880e-03,  3.1129e-03,  2.3662e-03,  6.9493e-05,\n        -2.8218e-03, -2.0927e-03,  2.6526e-04, -2.2561e-03,  1.5423e-03,\n        -1.3247e-04, -1.4266e-03], device='cuda:0')",
-        "exp_avg_sq": "tensor([7.2805e-05, 1.1672e-04, 1.0106e-04, 5.2318e-05, 6.4069e-05, 9.7167e-05,\n        6.2720e-05, 4.8752e-05, 6.3069e-05, 6.1414e-05, 6.9519e-05, 1.5512e-05,\n        1.0579e-04, 1.3858e-04, 4.9242e-05, 1.0395e-04, 7.9582e-05, 6.5475e-05,\n        6.4109e-05, 8.9945e-05, 6.4051e-05, 7.3161e-05, 1.2849e-04, 4.9342e-05,\n        1.2649e-04, 7.0120e-05, 7.6571e-05, 7.1208e-05, 5.9177e-05, 5.9676e-05,\n        9.4032e-05, 7.5547e-05, 6.9886e-05, 6.3639e-05, 1.4910e-04, 5.1891e-05,\n        1.4909e-04, 7.8397e-05, 4.3016e-05, 7.9174e-05, 7.4907e-05, 7.7912e-05,\n        8.4969e-05, 6.4953e-05, 7.6125e-05, 6.5005e-05, 8.3272e-05, 4.1785e-05,\n        4.3865e-05, 7.1190e-05, 2.3132e-05, 7.8169e-05, 5.0765e-05, 1.3006e-04,\n        6.5071e-05, 6.8094e-05, 7.3109e-05, 5.8487e-05, 4.6976e-05, 9.6623e-05,\n        1.1592e-04, 8.2575e-05, 7.3914e-05, 2.6722e-05, 6.6206e-05, 6.6949e-05,\n        7.3332e-05, 8.7034e-05, 7.6856e-05, 8.6299e-05, 1.1663e-04, 9.5386e-05,\n        7.4992e-05, 6.9103e-05, 5.6621e-05, 1.0645e-04, 8.0133e-05, 6.7468e-05,\n        5.5392e-05, 9.5648e-05, 9.7901e-05, 5.8464e-05, 5.9976e-05, 8.0588e-05,\n        5.6626e-05, 1.0551e-04, 1.1118e-04, 1.3983e-04, 1.9610e-05, 6.7562e-05,\n        1.2050e-04, 1.4105e-04, 9.6316e-05, 5.9347e-05, 7.3845e-05, 5.2278e-05,\n        8.9505e-05, 2.9763e-05, 6.9102e-05, 9.1485e-05, 4.9016e-05, 5.0686e-05,\n        6.8854e-05, 7.7402e-05, 7.1189e-05, 1.1991e-04, 9.6901e-05, 6.5220e-05,\n        8.6817e-05, 5.9262e-05, 5.5397e-05, 7.1925e-05, 6.0268e-05, 1.3699e-04,\n        5.6205e-05, 7.6740e-05, 8.1266e-05, 5.9240e-05, 6.3727e-05, 5.7301e-05,\n        6.8217e-05, 1.0202e-04, 7.3016e-05, 5.1731e-05, 5.4435e-05, 6.7459e-05,\n        6.2193e-05, 6.7981e-05, 7.1668e-05, 4.9179e-05, 6.6285e-05, 9.7299e-05,\n        9.1720e-05, 6.4078e-05, 5.5871e-05, 1.1454e-04, 8.2182e-05, 8.1519e-05,\n        6.6296e-05, 7.0125e-05, 9.6614e-05, 6.6239e-05, 8.0599e-05, 8.0127e-05,\n        7.6381e-05, 1.1804e-04, 9.0971e-05, 7.5930e-05, 9.8623e-05, 8.7635e-05,\n        7.9626e-05, 4.7248e-05, 6.6319e-05, 4.8574e-05, 8.5434e-05, 8.2103e-05,\n        8.6586e-05, 5.9175e-05, 1.4231e-04, 9.7886e-05, 9.6988e-05, 4.9253e-05,\n        6.2231e-05, 9.9259e-05, 3.3322e-05, 5.9274e-05, 1.0413e-04, 7.1749e-05,\n        7.4594e-05, 7.7133e-05, 6.0479e-05, 9.0234e-05, 7.5887e-05, 1.0110e-04,\n        7.4340e-05, 9.1068e-05, 8.9490e-05, 5.3439e-05, 4.7686e-05, 1.0250e-04,\n        1.8285e-04, 8.4924e-05, 5.3520e-05, 1.1323e-04, 6.1649e-05, 5.6565e-05,\n        6.7477e-05, 6.7504e-05, 5.8334e-05, 6.9376e-05, 7.3840e-05, 1.2862e-04,\n        8.0710e-05, 6.9240e-05, 9.3052e-05, 8.3365e-05, 6.2373e-05, 5.4119e-05,\n        5.9791e-05, 6.5900e-05, 8.2064e-05, 6.3508e-05, 5.4196e-05, 1.1574e-04,\n        4.7134e-05, 4.3102e-05, 6.5892e-05, 5.6867e-05, 6.7448e-05, 9.6472e-05,\n        6.1837e-05, 6.7980e-05, 8.4563e-05, 5.1236e-05, 1.5447e-04, 6.2577e-05,\n        7.2474e-05, 7.0674e-05, 5.0201e-05, 7.0031e-05, 8.6464e-05, 1.1326e-04,\n        9.8783e-05, 5.8473e-05, 8.8342e-05, 9.1716e-05, 9.3797e-05, 9.2950e-05,\n        6.7491e-05, 8.0675e-05, 6.1400e-05, 9.7006e-05, 6.0568e-05, 5.7513e-05,\n        4.6434e-05, 7.2067e-05, 1.0037e-04, 1.1206e-04, 9.2204e-05, 6.8287e-05,\n        7.4306e-05, 6.3577e-05, 6.5473e-05, 5.7117e-05, 5.3415e-05, 1.2839e-04,\n        6.8505e-05, 8.9959e-05, 5.2558e-05, 6.0559e-05, 5.7457e-05, 5.1058e-05,\n        8.6491e-05, 6.0221e-05, 7.6424e-05, 7.4634e-05, 7.2953e-05, 1.2200e-04,\n        7.6044e-05, 8.3126e-05, 8.2038e-05, 6.6028e-05, 1.2521e-04, 6.2375e-05,\n        6.2980e-05, 7.4432e-05, 5.8902e-05, 6.4820e-05, 4.7210e-05, 1.0792e-04,\n        4.3893e-05, 5.8598e-05, 8.6191e-05, 5.3834e-05, 7.0819e-05, 7.2045e-05,\n        7.2608e-05, 4.8905e-05, 8.3287e-05, 6.4034e-05, 1.1725e-04, 7.5680e-05,\n        8.2510e-05, 9.3443e-05, 6.3192e-05, 1.1004e-04, 4.4910e-05, 8.6279e-05,\n        6.2944e-05, 3.6750e-05, 8.2400e-05, 2.1229e-04, 8.0459e-05, 6.7904e-05,\n        8.1323e-05, 5.1442e-05, 1.4593e-04, 1.3080e-04, 8.6566e-05, 8.6640e-05,\n        5.5147e-05, 6.6972e-05, 8.1071e-05, 1.0103e-04, 1.4062e-04, 8.1013e-05,\n        8.1224e-05, 4.1284e-05, 6.3023e-05, 6.7015e-05, 5.6340e-05, 6.6877e-05,\n        5.8843e-05, 8.4409e-05, 6.2503e-05, 5.7634e-05, 4.7090e-05, 5.7766e-05,\n        8.1777e-05, 8.9744e-05, 5.4830e-05, 1.1836e-04, 1.4961e-04, 1.5467e-04,\n        4.8640e-05, 5.7058e-05, 5.9751e-05, 8.9172e-05, 4.1980e-05, 7.0580e-05,\n        5.8323e-05, 8.6059e-05, 8.6024e-05, 6.7151e-05, 5.7660e-05, 6.5126e-05,\n        8.2894e-05, 7.8449e-05, 1.1272e-04, 9.5756e-05, 6.0674e-05, 6.5323e-05,\n        5.5882e-05, 7.8808e-05, 8.5638e-05, 9.1780e-05, 6.1911e-05, 7.9946e-05,\n        7.3067e-05, 9.9334e-05, 6.1605e-05, 6.3694e-05, 6.5771e-05, 9.3156e-05,\n        8.8675e-05, 6.6073e-05, 4.8191e-05, 8.5816e-05, 8.2800e-05, 6.4953e-05,\n        6.6536e-05, 5.0825e-05, 1.0631e-04, 6.1941e-05, 7.5813e-05, 8.8358e-05,\n        7.7688e-05, 1.1930e-04, 5.4968e-05, 3.9239e-05, 1.4051e-04, 7.1852e-05,\n        9.6921e-05, 3.3407e-05, 1.2215e-04, 1.7351e-04, 7.9325e-05, 6.0431e-05,\n        7.7842e-05, 4.0504e-05, 7.3966e-05, 7.0299e-05, 1.8133e-04, 8.2598e-05,\n        9.8152e-05, 5.1160e-05, 5.2500e-05, 9.4402e-05, 1.0552e-04, 1.5736e-04,\n        8.8692e-05, 4.9003e-05, 8.4667e-05, 4.9623e-05, 6.3787e-05, 8.3243e-05,\n        9.3210e-05, 5.8036e-05, 5.8227e-05, 9.6982e-05, 1.5625e-04, 1.3779e-04,\n        6.8637e-05, 6.2400e-05, 7.3630e-05, 1.1814e-04, 6.2708e-05, 3.7354e-05,\n        8.3551e-05, 5.8398e-05, 8.2951e-05, 7.7685e-05, 7.8463e-05, 6.1756e-05,\n        8.4898e-05, 5.1412e-05, 9.8902e-05, 1.2384e-04, 5.3099e-05, 1.1431e-04,\n        7.1459e-05, 9.0733e-05, 8.4459e-05, 8.8195e-05, 7.2895e-05, 7.8504e-05,\n        7.1499e-05, 8.8813e-05, 8.4379e-05, 1.0651e-04, 8.8514e-05, 7.0312e-05,\n        8.5652e-05, 4.4603e-05, 6.2909e-05, 6.5082e-05, 1.1047e-04, 6.2961e-05,\n        7.2271e-05, 6.3915e-05, 1.1729e-04, 7.7635e-05, 9.6146e-05, 6.1956e-05,\n        7.2829e-05, 7.1319e-05, 1.0299e-04, 5.2079e-05, 7.6336e-05, 4.7313e-05,\n        7.5666e-05, 6.1152e-05, 7.7616e-05, 5.3362e-05, 7.5697e-05, 7.3949e-05,\n        4.7809e-05, 7.1756e-05, 6.4426e-05, 5.7166e-05, 1.0147e-04, 4.7214e-05,\n        5.9636e-05, 8.1167e-05, 6.6442e-05, 7.1525e-05, 1.2103e-04, 1.3463e-04,\n        1.1842e-04, 8.1263e-05, 1.1432e-04, 3.8813e-05, 5.8315e-05, 7.0469e-05,\n        7.0583e-05, 5.7181e-05, 5.6708e-05, 6.8827e-05, 1.1962e-04, 7.0168e-05,\n        1.0042e-04, 1.4811e-04, 6.4259e-05, 6.5131e-05, 8.8050e-05, 1.3623e-04,\n        1.4329e-04, 9.1338e-05, 1.3071e-04, 6.2032e-05, 1.2297e-04, 7.2070e-05,\n        5.8367e-05, 8.6389e-05, 8.6888e-05, 9.6991e-05, 5.1231e-05, 5.4932e-05,\n        7.0490e-05, 3.5919e-05, 6.3765e-05, 9.2558e-05, 5.3069e-05, 8.8184e-05,\n        7.9905e-05, 6.6059e-05, 8.3056e-05, 7.1206e-05, 1.0307e-04, 6.6920e-05,\n        8.9094e-05, 7.7502e-05], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([ 6.3563e-03,  2.5102e-03, -2.3939e-03,  6.9390e-04, -5.5107e-03,\n         1.8048e-03, -1.3423e-03,  1.2744e-03,  1.2543e-03,  2.8361e-03,\n        -1.7246e-03,  2.1835e-03, -1.4989e-03, -3.1949e-03, -1.2977e-03,\n         2.1179e-03, -5.9291e-03, -9.2005e-04,  2.8861e-03,  3.0686e-03,\n         3.6523e-03, -7.2715e-04,  4.4408e-03,  2.2455e-03, -2.2390e-03,\n        -3.4705e-04, -1.1323e-04, -2.6541e-03, -9.0658e-05, -2.0350e-03,\n         6.4861e-04,  5.6058e-05,  4.3530e-04,  3.0172e-03, -1.1393e-03,\n         1.5743e-03, -1.9657e-03,  6.6744e-03, -6.2427e-04, -6.2990e-03,\n        -6.6796e-03,  2.2028e-04, -7.6008e-04, -1.3117e-03, -6.5553e-04,\n         2.9044e-03,  3.5508e-03,  1.5549e-03,  2.6519e-04,  1.7346e-03,\n         3.6852e-04,  5.0304e-04,  1.9452e-04,  7.0303e-04,  2.2807e-03,\n         4.2914e-03,  1.0266e-03,  2.2294e-03, -1.1829e-03,  1.3515e-03,\n         2.0346e-03, -2.9442e-04, -3.9549e-04,  2.2319e-03, -2.4483e-04,\n        -2.5583e-03,  4.4849e-03,  1.4333e-03, -4.4578e-03, -2.0497e-03,\n         1.1495e-02,  3.5215e-03,  1.2701e-04, -9.1266e-04,  6.6057e-04,\n        -2.2992e-03,  3.9980e-03, -7.1079e-05,  1.1422e-03, -2.3329e-03,\n         3.9809e-03, -6.3117e-04, -4.7632e-03, -2.8593e-03,  5.2501e-03,\n        -9.9090e-03, -8.7446e-05, -2.1033e-03,  1.3585e-03, -7.1950e-04,\n         1.5056e-03,  2.0026e-03,  2.2293e-03, -2.6758e-03, -9.9534e-04,\n         4.4668e-03,  1.8737e-03,  1.8360e-03,  3.8190e-03,  1.9203e-04,\n        -4.6890e-03,  1.8226e-03,  1.8874e-03,  3.6358e-03,  1.3807e-03,\n        -7.7577e-04, -7.9496e-04,  1.3677e-04,  8.1841e-04, -3.1814e-03,\n         1.8048e-03,  3.3553e-03, -4.9102e-03,  1.3808e-03,  5.1207e-03,\n         1.6721e-03, -4.7491e-03, -2.6198e-03,  2.8556e-03,  1.8389e-03,\n        -1.3770e-05, -5.7796e-03,  1.0991e-03,  3.1372e-03,  4.8819e-03,\n        -5.4394e-04, -1.4880e-03, -7.2622e-03, -2.2927e-04,  4.5269e-04,\n        -1.1851e-04, -4.7925e-03,  4.2491e-03, -2.0320e-04, -6.5289e-04,\n        -1.1051e-03, -1.5166e-03, -2.3110e-03,  1.2267e-03,  2.2656e-03,\n        -1.2047e-03,  3.9875e-03,  1.6500e-03, -2.7301e-03, -1.6858e-03,\n         2.3919e-03,  3.4813e-03,  1.6305e-03,  2.4193e-03,  2.3856e-03,\n        -7.6614e-04,  2.1073e-03, -7.2462e-03, -1.5324e-03, -3.6564e-03,\n         3.1190e-03,  4.2763e-03,  7.9173e-03,  8.0900e-04, -1.2533e-03,\n        -9.9852e-03, -7.8643e-04, -1.5498e-03,  1.8211e-03, -1.4373e-03,\n        -3.4651e-04,  1.3923e-03, -9.1538e-04, -2.9207e-04, -1.1321e-03,\n         9.9711e-04,  2.4242e-03,  1.2174e-03, -6.6163e-04, -1.1129e-03,\n         4.4727e-03,  3.5576e-05,  1.1692e-03,  1.4891e-03,  7.9591e-04,\n         4.2842e-03,  2.2273e-03, -1.0673e-02,  3.6652e-03, -1.1921e-03,\n         5.6809e-04, -2.4059e-04, -1.0264e-03, -1.3917e-04, -3.2738e-04,\n         6.8789e-04, -1.4678e-02,  1.9791e-03,  4.1252e-03,  6.2255e-04,\n        -1.4504e-03, -1.8916e-03,  2.2912e-03,  2.5363e-04,  1.0106e-03,\n        -2.3781e-04, -5.7982e-03, -6.7601e-05,  2.1338e-03,  1.2260e-03,\n         8.8970e-04,  9.6591e-04,  1.3678e-03, -1.1165e-03,  1.3344e-03,\n         7.4076e-04,  1.4390e-03, -2.3686e-03, -5.5181e-03, -1.3431e-03,\n         6.3643e-04, -6.6143e-03,  1.6457e-03, -4.4675e-04,  3.8409e-04,\n        -2.1455e-03, -8.0203e-04,  4.7499e-04,  7.5643e-04,  2.8416e-03,\n         1.2082e-03, -2.3145e-03, -4.4622e-03,  9.9061e-04, -5.4417e-04,\n        -1.5471e-03, -2.9405e-03,  7.3630e-04,  4.1378e-03, -4.1548e-03,\n         4.9746e-03,  3.6228e-03,  3.5885e-03,  8.3581e-04,  1.0634e-03,\n        -4.0015e-03,  2.1256e-03, -1.1948e-03,  6.5337e-04, -4.4733e-04,\n        -5.3818e-03,  1.0360e-03, -4.6869e-03,  1.4853e-03,  1.2977e-03,\n         1.3324e-03, -3.2049e-03,  2.3267e-03,  5.8186e-04,  2.6390e-03,\n        -6.4317e-04,  1.1045e-03,  3.9773e-03,  1.3893e-03,  1.8098e-03,\n        -4.4831e-03, -1.4275e-04,  4.2212e-03, -6.8254e-05, -1.2112e-04,\n        -7.4728e-03, -4.4847e-03,  1.6913e-03, -2.3447e-03, -2.3030e-03,\n        -3.8929e-04,  1.3200e-03, -4.6651e-03, -1.2995e-03, -8.0070e-04,\n         2.1448e-03,  1.0197e-03,  3.9658e-04,  1.1312e-03,  4.3220e-03,\n         1.2821e-03, -5.3671e-03,  2.4647e-03, -8.2521e-04,  1.5489e-03,\n         4.2633e-03, -1.3165e-03,  3.1668e-03,  2.0352e-03, -1.0567e-03,\n         5.2512e-03, -7.3721e-03, -1.3819e-03, -3.1426e-04,  2.8566e-03,\n         1.9785e-03, -2.3924e-03, -1.6155e-03,  2.4685e-03, -9.6952e-05,\n         1.0270e-03, -4.7090e-03,  4.4791e-03,  1.5376e-05,  5.2302e-04,\n        -9.0432e-04,  1.0867e-03, -1.2149e-03,  4.6787e-04, -4.2822e-03,\n        -5.6963e-03,  7.2298e-04, -2.1089e-03, -2.5055e-03, -2.6269e-03,\n        -1.9026e-03,  1.1965e-03, -2.7991e-03, -1.5413e-03,  1.5862e-03,\n        -3.0942e-04, -1.0107e-03, -4.9279e-04, -3.0776e-03,  2.7953e-03,\n        -6.1656e-03, -7.3021e-04, -4.3944e-03, -3.6002e-04,  1.9508e-04,\n        -1.0083e-03,  1.0220e-03,  2.3339e-03, -2.4182e-03,  4.1325e-04,\n         1.8667e-03,  3.6274e-03,  1.1361e-03, -6.2754e-04, -1.9324e-04,\n        -3.6546e-04,  3.1481e-04,  3.7374e-03,  1.3475e-03,  2.9976e-03,\n        -7.3433e-05,  2.9237e-03, -2.1311e-03, -7.1621e-03,  9.0230e-05,\n         4.4070e-03, -2.4974e-03,  2.1867e-03,  1.6351e-03,  9.3232e-04,\n         1.6370e-03, -2.1578e-03, -6.9997e-04, -9.2883e-04,  9.8079e-04,\n        -8.8313e-04,  6.4415e-04,  1.3825e-03, -5.7022e-04, -1.4116e-03,\n         1.3109e-03, -3.5973e-04, -9.7172e-03,  2.9686e-04, -1.5405e-03,\n        -8.0118e-04,  4.2690e-03, -1.8635e-03, -4.1665e-04,  3.4896e-03,\n         1.3644e-03,  3.2192e-03,  4.5342e-03, -4.4682e-03,  1.7362e-03,\n         1.7646e-03,  5.7927e-05, -6.0777e-03,  4.1298e-04, -3.8661e-04,\n        -4.4611e-04,  3.9875e-04, -6.9911e-04,  4.4768e-03,  1.8348e-04,\n         3.3966e-03,  2.6140e-03, -8.2243e-03, -6.8521e-04, -2.9733e-04,\n         4.7474e-03, -3.4884e-03,  4.5268e-03, -8.1612e-04, -7.6521e-04,\n         2.2630e-03,  2.9607e-03,  2.9986e-03,  2.7078e-03, -2.5313e-03,\n        -5.4122e-03,  2.9183e-03, -1.5998e-03, -3.4358e-03, -2.9575e-04,\n        -2.0512e-03,  1.3428e-03,  1.0281e-03, -1.3787e-03,  3.1142e-03,\n        -7.7688e-04, -2.9966e-05, -2.5548e-03,  2.4785e-03,  1.2744e-04,\n         9.5441e-04, -2.6103e-03,  1.7985e-03,  3.4749e-03, -8.4321e-08,\n         1.6383e-03, -2.7722e-04,  6.5776e-03,  2.4065e-04,  3.0879e-03,\n        -7.4350e-03, -8.4307e-04, -7.2875e-04,  1.0637e-03, -2.5961e-03,\n         8.5130e-04, -1.6339e-03,  4.0686e-04,  6.7186e-04,  5.8664e-04,\n        -5.3609e-03, -1.3997e-04,  3.3515e-03, -8.5538e-04,  1.2525e-03,\n         8.4309e-04, -4.8851e-03,  1.0791e-03,  1.6683e-03,  1.8069e-03,\n         6.6044e-04,  1.4694e-03, -1.7793e-03, -3.6584e-05,  2.1455e-03,\n        -5.9628e-05, -3.8228e-03, -8.0960e-05, -2.9003e-03, -1.4925e-03,\n         1.1623e-03,  5.0000e-04,  1.6481e-03, -6.5982e-03,  3.4492e-04,\n        -2.4010e-03,  9.0007e-05, -1.1398e-04, -9.2671e-04, -1.2297e-03,\n         1.7663e-03, -1.6166e-03, -4.4402e-03,  1.3900e-04, -6.9665e-03,\n        -3.2943e-03, -1.8867e-03, -3.2147e-03,  9.4666e-04, -9.1608e-04,\n         2.3092e-03,  9.5321e-04, -1.8345e-03,  4.1769e-04,  1.1811e-03,\n        -1.9273e-03,  8.5545e-03, -6.6741e-04, -3.4244e-03,  3.2111e-03,\n         9.0464e-04,  1.4275e-03,  1.1919e-03, -3.3361e-03, -1.6580e-03,\n        -7.3152e-03, -5.0109e-04,  3.6467e-03,  6.5521e-07,  3.3706e-03,\n        -3.8867e-04, -4.0687e-04,  2.5767e-03, -1.3249e-03, -5.5586e-04,\n        -2.2303e-03,  5.3057e-03,  2.2691e-03, -4.4590e-03,  1.3507e-03,\n        -7.4353e-04,  5.9935e-03], device='cuda:0')",
+        "exp_avg_sq": "tensor([7.2512e-05, 1.1019e-04, 9.9248e-05, 5.5612e-05, 6.7121e-05, 1.0031e-04,\n        5.9773e-05, 5.1705e-05, 6.4033e-05, 5.6169e-05, 6.8238e-05, 1.6573e-05,\n        1.0072e-04, 1.4606e-04, 5.2493e-05, 1.0045e-04, 8.4599e-05, 6.0681e-05,\n        6.1493e-05, 9.0432e-05, 6.7030e-05, 7.0840e-05, 1.1957e-04, 4.8057e-05,\n        1.4041e-04, 7.1113e-05, 7.8978e-05, 6.9880e-05, 5.9417e-05, 5.7770e-05,\n        8.7513e-05, 7.1569e-05, 6.7041e-05, 6.2300e-05, 1.3242e-04, 5.0675e-05,\n        1.4801e-04, 7.9039e-05, 4.1354e-05, 7.4328e-05, 7.0986e-05, 7.9117e-05,\n        8.5124e-05, 6.4362e-05, 7.4290e-05, 6.6766e-05, 8.6604e-05, 4.0079e-05,\n        4.4369e-05, 7.3125e-05, 2.3740e-05, 7.4951e-05, 5.3456e-05, 1.2581e-04,\n        6.3913e-05, 7.1525e-05, 6.8563e-05, 5.9110e-05, 4.7262e-05, 1.0414e-04,\n        1.0640e-04, 8.7026e-05, 6.8453e-05, 2.7108e-05, 6.4277e-05, 6.4376e-05,\n        7.4997e-05, 9.0205e-05, 8.2962e-05, 8.7746e-05, 1.1170e-04, 9.4883e-05,\n        7.4454e-05, 6.6858e-05, 5.9579e-05, 1.0697e-04, 8.3371e-05, 6.6815e-05,\n        5.7789e-05, 9.6536e-05, 1.0008e-04, 6.2869e-05, 5.7528e-05, 8.2585e-05,\n        5.8706e-05, 1.0831e-04, 1.1401e-04, 1.3857e-04, 1.9935e-05, 6.5200e-05,\n        1.0378e-04, 1.3060e-04, 9.3773e-05, 5.6892e-05, 6.7640e-05, 4.9256e-05,\n        9.2273e-05, 2.9546e-05, 6.7948e-05, 9.2812e-05, 4.7165e-05, 4.8418e-05,\n        6.8732e-05, 8.0735e-05, 7.0170e-05, 1.1922e-04, 9.8124e-05, 6.6470e-05,\n        7.6156e-05, 5.9479e-05, 5.6384e-05, 7.0079e-05, 5.9407e-05, 1.3077e-04,\n        5.7598e-05, 7.4382e-05, 8.1519e-05, 5.9187e-05, 6.5026e-05, 5.9065e-05,\n        6.2860e-05, 1.0174e-04, 7.1195e-05, 5.0764e-05, 5.7857e-05, 6.4192e-05,\n        6.1678e-05, 7.2441e-05, 7.3560e-05, 4.9425e-05, 6.4231e-05, 9.7638e-05,\n        9.1556e-05, 6.4483e-05, 5.9572e-05, 1.0801e-04, 8.4543e-05, 8.1038e-05,\n        6.6712e-05, 6.6604e-05, 9.2933e-05, 6.4318e-05, 8.5129e-05, 7.6727e-05,\n        7.6444e-05, 1.0942e-04, 8.7810e-05, 7.8066e-05, 9.8640e-05, 8.4970e-05,\n        7.0386e-05, 4.5215e-05, 6.9736e-05, 4.8611e-05, 8.0937e-05, 7.8415e-05,\n        9.6527e-05, 6.2934e-05, 1.2635e-04, 9.8119e-05, 1.0523e-04, 5.4049e-05,\n        6.6602e-05, 9.7183e-05, 3.0206e-05, 5.5845e-05, 1.0028e-04, 6.6852e-05,\n        7.3802e-05, 7.6700e-05, 5.7336e-05, 9.6055e-05, 7.6344e-05, 1.0505e-04,\n        7.6953e-05, 8.1130e-05, 8.5807e-05, 5.1258e-05, 4.4070e-05, 9.9288e-05,\n        1.8838e-04, 8.6845e-05, 6.1611e-05, 1.0970e-04, 6.3698e-05, 5.3190e-05,\n        7.1777e-05, 6.9004e-05, 5.6626e-05, 7.1177e-05, 6.9374e-05, 1.4085e-04,\n        8.0837e-05, 6.9333e-05, 8.4503e-05, 8.3074e-05, 6.1865e-05, 5.7168e-05,\n        6.1964e-05, 6.7875e-05, 7.5165e-05, 6.0273e-05, 6.2879e-05, 1.1144e-04,\n        5.1099e-05, 4.3593e-05, 6.0414e-05, 5.3644e-05, 6.1797e-05, 9.7747e-05,\n        6.1278e-05, 7.1418e-05, 8.3301e-05, 5.3395e-05, 1.4802e-04, 6.1757e-05,\n        7.9729e-05, 7.4921e-05, 4.8466e-05, 6.7437e-05, 8.6289e-05, 1.1120e-04,\n        1.0087e-04, 5.7920e-05, 8.3983e-05, 8.7964e-05, 9.8691e-05, 8.5369e-05,\n        6.4883e-05, 7.9041e-05, 6.0690e-05, 9.2641e-05, 6.0072e-05, 5.4724e-05,\n        4.6984e-05, 7.3237e-05, 9.9077e-05, 1.1399e-04, 8.2742e-05, 7.2863e-05,\n        7.8897e-05, 5.9803e-05, 6.4471e-05, 5.7344e-05, 5.6070e-05, 1.3041e-04,\n        6.6416e-05, 8.8161e-05, 5.2970e-05, 6.4837e-05, 5.6148e-05, 5.0703e-05,\n        8.3795e-05, 5.7469e-05, 8.0018e-05, 7.3807e-05, 7.5782e-05, 1.1732e-04,\n        7.1639e-05, 8.2258e-05, 7.8587e-05, 6.4899e-05, 1.2533e-04, 5.7355e-05,\n        6.2201e-05, 7.6339e-05, 6.1350e-05, 6.5109e-05, 4.5595e-05, 1.1440e-04,\n        4.2570e-05, 5.6618e-05, 8.4310e-05, 5.3958e-05, 6.7097e-05, 7.2834e-05,\n        6.8835e-05, 4.9115e-05, 8.2634e-05, 6.2827e-05, 1.1426e-04, 8.8805e-05,\n        8.5880e-05, 9.3828e-05, 6.4395e-05, 1.1473e-04, 4.3674e-05, 8.1672e-05,\n        6.1591e-05, 3.6387e-05, 8.3369e-05, 2.0206e-04, 7.8578e-05, 6.4915e-05,\n        7.8079e-05, 5.1374e-05, 1.3915e-04, 1.2021e-04, 8.2560e-05, 8.5514e-05,\n        5.4398e-05, 7.0501e-05, 8.9104e-05, 1.0418e-04, 1.2797e-04, 7.5903e-05,\n        7.4508e-05, 4.2443e-05, 6.1746e-05, 6.7860e-05, 5.8586e-05, 6.9703e-05,\n        5.9039e-05, 8.4074e-05, 6.1156e-05, 5.9695e-05, 4.6840e-05, 5.5527e-05,\n        8.3250e-05, 8.8701e-05, 5.5929e-05, 1.1869e-04, 1.6266e-04, 1.5948e-04,\n        4.7998e-05, 6.0044e-05, 5.6220e-05, 8.8606e-05, 4.1672e-05, 6.8777e-05,\n        5.7954e-05, 8.8043e-05, 7.8968e-05, 6.7892e-05, 5.7282e-05, 6.2787e-05,\n        7.5525e-05, 8.0414e-05, 1.1319e-04, 9.2808e-05, 5.8435e-05, 6.5216e-05,\n        5.4113e-05, 8.3818e-05, 8.2172e-05, 9.7351e-05, 6.6038e-05, 7.9793e-05,\n        8.1784e-05, 9.9960e-05, 6.0991e-05, 6.2078e-05, 6.4862e-05, 8.8328e-05,\n        7.8428e-05, 6.7388e-05, 4.5181e-05, 8.6514e-05, 7.8928e-05, 6.4784e-05,\n        6.2382e-05, 4.9389e-05, 9.9456e-05, 6.1233e-05, 7.9262e-05, 9.6575e-05,\n        7.0088e-05, 1.2153e-04, 5.3392e-05, 3.9271e-05, 1.4991e-04, 6.4224e-05,\n        9.3520e-05, 3.2669e-05, 1.1926e-04, 1.6915e-04, 7.5040e-05, 6.2417e-05,\n        8.0370e-05, 3.7825e-05, 7.3172e-05, 6.6714e-05, 1.7721e-04, 8.2638e-05,\n        1.1069e-04, 5.3679e-05, 5.0091e-05, 9.0835e-05, 1.0999e-04, 1.5140e-04,\n        8.8529e-05, 5.0547e-05, 8.3844e-05, 4.8599e-05, 6.3352e-05, 7.7482e-05,\n        9.3653e-05, 5.5084e-05, 5.7589e-05, 9.5163e-05, 1.5871e-04, 1.2744e-04,\n        6.9541e-05, 6.4151e-05, 7.0628e-05, 1.3157e-04, 6.4501e-05, 3.7270e-05,\n        8.6574e-05, 5.7349e-05, 8.2512e-05, 7.2937e-05, 7.7274e-05, 5.7988e-05,\n        8.3994e-05, 4.8179e-05, 9.4418e-05, 1.1850e-04, 5.1639e-05, 1.2304e-04,\n        7.3215e-05, 8.6787e-05, 8.2314e-05, 8.7038e-05, 7.4435e-05, 7.7177e-05,\n        7.2046e-05, 9.2039e-05, 8.2250e-05, 1.0159e-04, 8.9180e-05, 6.8441e-05,\n        8.5019e-05, 4.4265e-05, 6.2266e-05, 6.2828e-05, 1.0825e-04, 6.2090e-05,\n        6.6607e-05, 6.1238e-05, 1.1401e-04, 8.1602e-05, 9.3099e-05, 6.1156e-05,\n        7.3361e-05, 7.1485e-05, 9.7311e-05, 5.0915e-05, 7.6193e-05, 4.7340e-05,\n        7.3933e-05, 6.2318e-05, 7.6087e-05, 5.0242e-05, 7.3717e-05, 7.1194e-05,\n        5.5668e-05, 7.0081e-05, 6.7980e-05, 5.7673e-05, 1.0371e-04, 4.7619e-05,\n        5.7822e-05, 8.2032e-05, 6.2093e-05, 7.3440e-05, 1.2068e-04, 1.3620e-04,\n        1.1772e-04, 7.7936e-05, 1.1727e-04, 3.7691e-05, 5.7127e-05, 6.8972e-05,\n        6.2754e-05, 5.4889e-05, 5.8448e-05, 7.1129e-05, 1.1267e-04, 7.0763e-05,\n        9.7364e-05, 1.5761e-04, 6.3211e-05, 6.2182e-05, 8.9459e-05, 1.2892e-04,\n        1.4036e-04, 8.5110e-05, 1.2891e-04, 6.0253e-05, 1.0926e-04, 7.3221e-05,\n        5.9762e-05, 8.5187e-05, 8.6531e-05, 1.0131e-04, 5.0183e-05, 5.5968e-05,\n        7.5582e-05, 3.8386e-05, 6.1099e-05, 9.0729e-05, 5.2793e-05, 9.0854e-05,\n        7.4200e-05, 6.2664e-05, 8.1053e-05, 6.8687e-05, 1.0426e-04, 6.5999e-05,\n        9.0164e-05, 7.7003e-05], device='cuda:0')"
       },
       "2": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([-2.1598e-03, -1.3331e-02,  2.4635e-03, -3.1663e-04, -9.2909e-03,\n        -1.2417e-03, -1.8650e-04,  4.0231e-03,  5.6357e-03, -5.7458e-03,\n        -4.4149e-03, -2.7567e-03,  7.0216e-03, -2.5503e-03,  1.9709e-04,\n        -1.7082e-03,  8.7501e-03,  1.6027e-03,  1.0806e-02, -3.7420e-03,\n        -2.0458e-03, -1.0274e-03,  1.1167e-03,  6.4120e-04, -6.3697e-03,\n        -9.6420e-04,  2.3035e-03, -1.4807e-03,  1.0856e-04, -1.7119e-02,\n        -5.7458e-04, -1.3814e-02,  1.6994e-03, -1.3046e-03, -3.8044e-03,\n         3.9457e-05, -4.5246e-04,  3.6382e-03, -8.0086e-04, -3.3888e-03,\n         1.2944e-03, -6.4929e-04, -4.6175e-03, -5.3734e-03, -3.8478e-03,\n        -2.9502e-03,  1.0236e-03,  8.2416e-03,  3.4605e-03,  2.6239e-04,\n        -5.9029e-03, -2.1931e-03, -1.5148e-03, -9.1489e-04, -2.1994e-03,\n         2.4297e-04,  6.5189e-03, -5.5951e-03,  9.3219e-03,  6.0254e-03,\n         3.1115e-03,  1.8645e-03, -7.3853e-03, -4.8424e-02,  4.6796e-03,\n         3.9203e-03, -4.4192e-03, -4.8143e-03, -8.0104e-04,  1.8322e-03,\n         1.1858e-03, -2.7953e-03, -9.4203e-03, -3.0664e-04,  2.2903e-03,\n        -8.5703e-04, -2.3320e-03,  1.7026e-03,  2.7189e-03, -3.7939e-04,\n        -9.1520e-03, -8.9197e-03, -8.6957e-03,  1.3533e-03,  1.7303e-04,\n        -3.0344e-04, -3.7479e-03,  4.3237e-03,  5.6052e-45,  3.3666e-03,\n        -6.5301e-03, -4.6150e-03,  3.0114e-04, -2.9108e-03,  1.6715e-02,\n        -2.2654e-03, -2.7004e-03,  6.6019e-03, -1.0103e-03, -8.2723e-03,\n         7.2052e-04, -2.0948e-03, -2.5141e-03, -3.7424e-03, -7.0037e-03,\n         2.6711e-03, -3.1580e-03,  4.6240e-03, -2.5422e-03, -9.7807e-04,\n         5.5428e-04, -1.7860e-03, -1.6757e-04,  3.4493e-05, -1.2321e-03,\n         7.2926e-03, -4.0637e-04, -3.5699e-04,  1.8565e-04, -6.7726e-04,\n         4.1459e-03,  4.1983e-03,  1.6943e-03,  2.6251e-03, -1.0065e-02,\n        -9.2742e-03, -8.3702e-04, -1.8104e-03,  3.4891e-03, -1.6025e-03,\n        -1.3280e-03, -3.6686e-03, -3.2964e-03, -6.7988e-04,  6.2835e-04,\n        -6.4328e-04,  1.3178e-03,  2.7811e-03, -9.3120e-04, -1.2135e-02,\n         3.0747e-03, -4.8724e-03, -2.4937e-03, -5.5626e-03, -1.9687e-03,\n        -3.1175e-03,  1.1417e-03, -2.6216e-03, -3.4784e-04, -5.1030e-03,\n        -4.7417e-03, -2.3809e-03, -5.8985e-04,  1.5073e-03, -1.9702e-03,\n        -4.9919e-03,  1.8621e-03,  2.0517e-03,  6.9460e-03,  2.3744e-03,\n         6.7827e-04, -2.7661e-03, -8.8280e-04, -6.8954e-04,  3.8842e-04,\n         4.0790e-03,  3.7194e-04,  2.2308e-03,  1.2398e-03,  6.0004e-04,\n         1.7130e-03,  3.5784e-03, -3.1007e-03,  8.4413e-04, -4.0617e-04,\n        -1.3512e-03,  6.3770e-03,  2.6323e-03,  8.0931e-03, -7.0408e-03,\n         3.3100e-03,  1.1808e-03, -2.5501e-03,  1.7726e-03,  1.3648e-03,\n        -2.4245e-03,  1.4567e-03, -3.0202e-03, -4.7279e-03,  3.3748e-03,\n         3.9464e-03, -2.2582e-03,  1.1225e-02, -9.6889e-04, -5.7278e-05,\n         2.0989e-03,  4.6267e-03, -1.3822e-03,  8.8851e-04,  2.5126e-03,\n         2.1377e-03, -1.1509e-02,  1.0137e-03, -8.7899e-03, -6.7968e-04,\n        -4.4979e-04, -2.2699e-03, -1.3911e-03,  1.3414e-03, -3.3662e-03,\n         4.5313e-04,  2.5125e-03,  5.7528e-03,  1.6893e-03,  2.9572e-03,\n         7.6680e-03,  2.7938e-03,  4.3962e-04,  2.0400e-03,  8.4747e-03,\n         3.4371e-03,  6.0308e-03,  5.9122e-04, -3.1993e-03,  3.3645e-03,\n         2.6794e-03,  6.5403e-04,  9.9614e-04,  3.2798e-03, -8.3289e-04,\n        -1.0815e-04, -1.2461e-04,  2.1699e-03,  4.1766e-03, -2.6223e-03,\n         1.4188e-03,  4.7749e-03, -1.4593e-03,  3.5088e-07, -1.5172e-03,\n         7.7114e-03, -3.3059e-03,  2.2752e-03, -2.8104e-03, -3.3559e-03,\n        -2.9414e-03, -9.8947e-04,  5.5870e-03,  4.0328e-04,  5.6586e-04,\n        -7.6475e-03, -1.8352e-03, -3.9980e-03,  5.0869e-04, -2.6624e-04,\n         2.4606e-03,  2.4681e-04, -8.2943e-03, -1.7594e-03,  1.3387e-03,\n        -5.7861e-05, -1.3539e-03,  2.3317e-03, -3.7471e-05,  1.3667e-03,\n         1.9595e-03,  1.4668e-03, -3.0699e-03, -6.5464e-04, -2.6079e-03,\n        -4.8374e-03,  1.6512e-03,  2.5132e-03, -2.4195e-04,  1.6259e-03,\n         1.0779e-02,  4.1774e-03, -2.0666e-03, -6.2359e-03, -1.9261e-03,\n         2.2068e-03,  3.1091e-03,  4.0737e-03, -3.0657e-04,  8.3369e-03,\n         6.2106e-03,  2.2040e-03,  6.0858e-04,  2.4888e-03,  1.3773e-03,\n         3.3892e-03,  1.7468e-03,  6.0913e-04, -1.5053e-04, -1.3756e-02,\n        -2.9547e-03, -4.1643e-03,  3.4925e-03,  4.0894e-03,  7.0869e-03,\n        -1.1883e-04,  8.2540e-03,  2.3744e-03, -5.5347e-03,  2.2038e-03,\n        -5.9507e-03, -9.9145e-03,  9.1675e-05, -1.6520e-03, -4.5656e-04,\n         2.0493e-03, -1.0919e-03, -5.9612e-04, -2.9907e-03,  1.9414e-03,\n        -1.8267e-04, -9.3254e-04,  2.1558e-03, -8.5991e-03, -3.9547e-04,\n        -1.5434e-03,  3.0403e-03, -7.0743e-04, -3.1864e-03,  8.2458e-04,\n         5.0944e-04, -5.6710e-03, -4.1944e-03, -2.6184e-04,  2.7893e-03,\n         6.9264e-03,  6.7511e-03, -4.8632e-04, -2.5416e-03,  2.6133e-03,\n        -5.3130e-03, -4.6601e-03, -2.2476e-03, -3.1639e-03, -5.9086e-05,\n        -1.3946e-03, -1.4716e-03,  4.7066e-03,  7.8387e-04,  7.4879e-04,\n        -8.2580e-03,  3.4215e-03, -7.4760e-04, -5.3507e-04,  6.1072e-03,\n         3.1610e-03, -3.3264e-03,  3.4175e-03, -2.8842e-03, -7.1520e-03,\n        -4.5539e-03, -6.3616e-03,  4.2525e-04, -6.0094e-03,  1.4183e-03,\n         1.0155e-02,  1.3742e-03, -5.1974e-04, -3.5589e-03,  3.8131e-04,\n        -7.4190e-04, -3.5671e-03,  1.8117e-02, -1.2142e-02,  4.1505e-04,\n        -3.9797e-03,  1.8902e-03, -6.4457e-04, -1.0129e-03, -4.9364e-03,\n         6.8166e-05,  4.6230e-04,  7.6259e-03,  4.5556e-03,  5.3753e-03,\n        -7.4658e-04, -3.1085e-03, -3.4856e-03,  4.3359e-03,  2.0977e-03,\n        -3.6066e-03,  9.6025e-04, -2.9420e-03,  1.9615e-03, -7.6115e-03,\n         2.8944e-03,  6.9764e-03,  4.2157e-03,  2.7888e-03, -3.6343e-03,\n        -5.7792e-03, -7.7173e-04, -5.3464e-03,  3.3089e-03, -6.5740e-03,\n        -2.9229e-03,  2.9329e-03,  1.5747e-03, -5.7693e-04, -1.7160e-03,\n         4.5562e-03, -1.2604e-04, -3.9503e-04,  2.5259e-03,  2.9154e-04,\n         9.0616e-04,  2.8032e-04,  3.4618e-03,  3.0292e-03,  2.1573e-03,\n         5.1566e-04,  4.5950e-03,  8.1460e-04,  1.7370e-02,  4.5326e-03,\n         4.6098e-04,  8.9486e-03, -2.1099e-03, -3.8010e-03,  2.4552e-03,\n         1.0853e-03,  4.6139e-04,  1.1698e-03,  5.2924e-03, -8.8034e-06,\n        -6.0701e-03, -7.8123e-04, -3.1566e-03,  1.0378e-02,  3.7926e-03,\n        -4.4217e-04, -1.8508e-02, -7.0008e-03,  1.9072e-03, -5.2951e-03,\n         2.6164e-03, -9.8213e-04,  3.2002e-03,  2.8028e-03, -1.7904e-03,\n        -1.3218e-03, -2.0728e-04,  1.8907e-03, -1.2114e-03,  4.4000e-03,\n         4.6128e-03, -1.4851e-03, -3.1792e-03,  2.7323e-03,  3.4561e-03,\n        -4.5618e-03,  4.3595e-03,  3.4601e-03, -3.5447e-03,  3.6863e-03,\n        -1.9330e-03, -1.1415e-04, -2.6319e-03, -3.6117e-03, -2.1152e-03,\n         1.2409e-03,  1.1122e-02,  3.6606e-03,  2.7383e-03,  1.2092e-02,\n        -6.9677e-03, -3.6998e-03, -1.2203e-03, -1.4037e-03,  4.6699e-03,\n         1.4130e-03,  2.9021e-03,  1.1085e-03, -1.1719e-04, -5.6215e-05,\n         2.9496e-03,  1.5727e-02,  3.8869e-03, -1.4298e-03,  4.4690e-03,\n        -2.0902e-03, -4.5195e-03, -2.0456e-03,  2.4114e-03,  2.4587e-03,\n         1.2828e-03,  1.8653e-03,  2.3077e-03,  4.1995e-03,  2.1422e-03,\n         2.3106e-03,  3.6190e-03, -1.7287e-07,  1.5162e-03, -1.9856e-05,\n        -4.0081e-03, -2.2682e-03,  3.6068e-03,  2.6536e-03, -4.5523e-04,\n        -2.5734e-03, -4.2842e-03,  1.1082e-03, -4.0712e-03,  2.6034e-03,\n         4.0521e-04, -3.2018e-03], device='cuda:0')",
-        "exp_avg_sq": "tensor([7.0456e-05, 1.5009e-04, 1.1914e-04, 1.1305e-04, 1.2948e-04, 9.3458e-05,\n        5.8144e-05, 8.9416e-05, 1.5834e-04, 1.5016e-04, 1.3475e-04, 1.0399e-03,\n        1.5464e-04, 1.4707e-04, 2.0417e-04, 7.4560e-05, 1.9788e-04, 2.4191e-04,\n        2.2815e-04, 1.1510e-04, 1.4176e-04, 8.9531e-05, 1.1135e-04, 7.0938e-05,\n        1.9299e-04, 8.5557e-05, 1.5959e-04, 1.3847e-04, 2.7688e-04, 2.8286e-04,\n        1.0896e-04, 1.6325e-04, 7.5875e-05, 6.8070e-05, 2.0709e-04, 1.0257e-04,\n        1.2006e-04, 1.9336e-04, 1.5027e-04, 9.5387e-05, 1.0976e-04, 1.1785e-04,\n        1.2511e-04, 1.4879e-04, 7.9793e-05, 1.0988e-04, 1.2627e-04, 1.8137e-04,\n        2.8450e-04, 7.6617e-05, 6.4563e-04, 3.1285e-04, 1.2545e-04, 8.6273e-05,\n        1.5786e-04, 1.0979e-04, 1.4674e-04, 1.8622e-04, 2.0073e-04, 2.0302e-04,\n        2.1901e-04, 1.2259e-04, 2.3305e-04, 6.7960e-03, 1.2712e-04, 1.3401e-04,\n        1.7447e-04, 1.4434e-04, 5.0623e-05, 2.0483e-04, 9.6911e-05, 1.9276e-04,\n        1.5472e-04, 1.0826e-04, 1.0037e-04, 2.8092e-04, 8.8502e-05, 3.8089e-04,\n        1.8838e-04, 1.5330e-04, 2.9146e-04, 4.6684e-04, 6.0674e-04, 1.2559e-04,\n        1.8947e-04, 8.1768e-05, 9.5319e-05, 1.8006e-04, 1.4973e-08, 2.4467e-04,\n        1.1484e-04, 2.3688e-04, 1.4467e-04, 1.5962e-04, 3.6312e-04, 1.5016e-04,\n        8.7344e-05, 4.7232e-04, 1.0972e-04, 2.4619e-04, 1.7908e-04, 9.7160e-05,\n        1.0756e-04, 3.3452e-04, 1.6882e-04, 7.8678e-05, 8.3191e-05, 9.0641e-05,\n        1.1862e-04, 1.8179e-04, 2.0383e-04, 1.1656e-04, 1.4504e-04, 6.9762e-05,\n        1.2412e-04, 1.1605e-04, 1.2001e-04, 1.2739e-04, 1.1002e-04, 5.1165e-05,\n        2.0757e-04, 1.0083e-04, 1.7286e-04, 2.0560e-04, 1.3855e-04, 9.3621e-05,\n        7.3077e-05, 7.1392e-05, 9.4990e-05, 2.0855e-04, 1.1957e-04, 1.8627e-04,\n        1.6050e-04, 2.4748e-04, 3.1321e-04, 2.0522e-04, 1.1371e-04, 1.2945e-04,\n        9.0055e-05, 2.3393e-04, 1.4301e-04, 2.0417e-04, 1.1801e-04, 1.6835e-04,\n        1.6346e-04, 1.5477e-04, 2.5933e-04, 1.3581e-04, 1.0254e-04, 1.6828e-04,\n        8.8327e-05, 1.7738e-04, 1.0771e-04, 1.5278e-04, 9.5033e-05, 1.4736e-04,\n        1.1269e-04, 5.4912e-05, 1.5553e-04, 1.1116e-04, 9.4713e-05, 9.8400e-05,\n        1.8346e-04, 1.0082e-04, 1.2781e-04, 1.8925e-04, 7.4123e-05, 2.3166e-04,\n        1.5416e-04, 2.1030e-04, 1.3197e-04, 1.8065e-04, 1.3764e-04, 1.1571e-04,\n        3.5938e-04, 1.3400e-04, 1.5570e-04, 2.6593e-04, 2.0473e-04, 2.1273e-04,\n        1.1368e-04, 1.5148e-04, 1.9944e-04, 4.3393e-04, 1.9495e-04, 1.0815e-04,\n        2.1662e-04, 3.8200e-04, 1.4464e-04, 1.8513e-04, 1.2005e-04, 2.1202e-04,\n        2.4040e-04, 7.4149e-05, 1.9140e-04, 7.8894e-05, 1.7043e-04, 1.0785e-04,\n        1.0631e-04, 7.5467e-05, 1.6808e-04, 2.9017e-04, 1.9578e-04, 1.4678e-04,\n        1.3161e-04, 1.9622e-04, 1.2016e-04, 1.1884e-04, 7.3653e-05, 9.3064e-05,\n        1.1042e-04, 2.6398e-04, 2.0247e-04, 1.0830e-04, 1.5303e-04, 1.4675e-04,\n        8.3286e-05, 2.0008e-04, 3.1766e-04, 1.6430e-04, 2.1337e-04, 1.3933e-04,\n        1.9404e-04, 1.6279e-04, 2.1034e-04, 5.6324e-05, 2.1700e-04, 1.0644e-04,\n        1.5092e-04, 1.0343e-04, 4.7726e-05, 1.1820e-04, 1.8846e-04, 2.4895e-04,\n        2.7210e-04, 1.3519e-04, 2.8500e-04, 2.2263e-04, 6.8412e-05, 8.5112e-05,\n        1.5189e-04, 7.9421e-05, 9.2352e-05, 1.8886e-04, 9.6438e-05, 1.1327e-04,\n        7.6249e-05, 1.3885e-04, 3.7313e-04, 1.3231e-04, 2.0939e-04, 1.7119e-04,\n        1.7016e-04, 1.1966e-04, 1.0787e-04, 1.5269e-04, 7.1132e-05, 2.4259e-04,\n        6.4263e-05, 1.2002e-04, 1.3578e-04, 2.4295e-04, 1.4125e-04, 9.7106e-05,\n        1.6376e-04, 2.3540e-04, 1.4530e-04, 7.9351e-05, 9.6709e-05, 1.8006e-04,\n        1.2402e-04, 9.5453e-05, 1.8862e-04, 1.0408e-04, 1.3135e-04, 4.7619e-04,\n        1.7293e-04, 1.4682e-04, 1.9560e-04, 2.2356e-04, 1.6138e-04, 1.3363e-04,\n        3.5804e-04, 1.2104e-04, 1.2106e-04, 1.4162e-04, 1.1349e-04, 2.9605e-04,\n        2.4199e-04, 2.5788e-04, 7.5528e-05, 1.9709e-04, 1.0232e-04, 7.4936e-05,\n        1.6419e-04, 1.6695e-04, 2.8572e-04, 3.6373e-04, 5.5839e-05, 3.0319e-04,\n        1.2576e-04, 6.6974e-04, 8.2076e-05, 7.4872e-05, 1.2305e-04, 2.4719e-04,\n        2.3909e-04, 2.0047e-04, 6.9104e-05, 1.2146e-04, 2.3945e-04, 1.1175e-04,\n        8.6150e-05, 1.3995e-04, 7.1974e-05, 8.2153e-05, 2.5249e-04, 2.2916e-04,\n        1.3448e-04, 1.2852e-04, 2.2778e-04, 2.6800e-04, 1.5432e-04, 1.4518e-04,\n        2.8454e-04, 9.1885e-05, 1.3150e-04, 1.1321e-04, 1.7023e-04, 2.0136e-04,\n        1.7274e-04, 2.1113e-04, 1.6879e-04, 1.1751e-04, 9.5951e-05, 1.2990e-04,\n        1.0664e-04, 1.6057e-04, 5.9900e-05, 1.9292e-04, 1.9632e-04, 1.4473e-04,\n        1.3624e-04, 1.7070e-04, 9.5692e-05, 2.1266e-04, 2.2991e-04, 1.4450e-04,\n        6.4448e-05, 3.2025e-04, 9.9368e-05, 2.4657e-04, 3.3304e-04, 8.1974e-05,\n        1.3996e-04, 2.5838e-04, 2.6441e-04, 1.3243e-04, 2.4719e-04, 8.3138e-05,\n        1.6314e-04, 7.8808e-05, 6.3365e-05, 1.7846e-04, 1.0606e-04, 3.6192e-05,\n        1.0272e-04, 3.6599e-04, 2.1993e-04, 1.2707e-04, 1.5204e-04, 1.5199e-04,\n        1.2041e-04, 1.5349e-04, 1.4014e-04, 3.7345e-04, 7.3700e-05, 3.6688e-04,\n        1.0481e-04, 2.3423e-04, 1.6598e-04, 1.3429e-04, 2.1209e-04, 2.6407e-04,\n        1.7780e-04, 3.1776e-04, 1.0886e-04, 1.4060e-04, 1.8186e-04, 1.8343e-04,\n        1.4500e-04, 4.2175e-04, 9.2360e-05, 1.8829e-04, 1.7036e-04, 1.4035e-04,\n        8.4034e-05, 2.6793e-04, 1.8107e-04, 6.0910e-05, 1.4992e-04, 1.3031e-04,\n        1.7300e-04, 9.6935e-05, 8.8861e-05, 2.7668e-04, 9.6777e-05, 1.9774e-04,\n        9.2582e-05, 1.6843e-04, 1.8386e-04, 1.8289e-04, 9.5230e-05, 1.2742e-04,\n        1.2209e-04, 1.5117e-04, 1.2388e-04, 1.8459e-04, 4.2954e-04, 1.5996e-04,\n        1.4289e-04, 3.0419e-04, 1.6153e-04, 1.8000e-04, 1.3778e-04, 1.4617e-04,\n        9.7568e-05, 1.7612e-04, 1.3531e-04, 8.4477e-05, 9.0790e-05, 1.3422e-04,\n        1.1834e-04, 6.8843e-04, 1.2080e-04, 1.9975e-04, 2.9919e-04, 1.3146e-04,\n        1.6635e-04, 2.5041e-04, 4.1398e-04, 9.4492e-05, 1.0834e-04, 2.1030e-04,\n        1.3784e-04, 1.9517e-04, 1.9101e-04, 1.7545e-04, 1.4157e-04, 2.6179e-04,\n        3.2544e-04, 1.2386e-04, 8.4592e-05, 2.9413e-04, 1.5362e-04, 2.5830e-04,\n        9.9989e-05, 2.4951e-04, 1.6470e-04, 1.4054e-04, 1.6963e-04, 3.3724e-04,\n        2.2494e-04, 1.9516e-04, 1.9142e-04, 4.0985e-04, 1.6539e-04, 1.7290e-04,\n        1.6193e-04, 1.6201e-04, 2.2907e-04, 1.1589e-04, 6.7867e-05, 1.4702e-04,\n        1.9824e-04, 1.2270e-04, 2.4364e-04, 1.7189e-04, 6.2261e-05, 1.1433e-04,\n        8.1606e-05, 1.9833e-04, 2.1331e-04, 6.9574e-05, 7.5029e-05, 1.0084e-04,\n        1.2801e-04, 1.2344e-04, 1.2637e-04, 8.5712e-05, 2.1129e-04, 1.6248e-04,\n        1.1741e-04, 2.0965e-04, 1.1904e-04, 1.5024e-04, 1.3089e-04, 1.0461e-04,\n        6.5309e-05, 3.5914e-05, 2.2506e-04, 7.6340e-05, 1.2981e-04, 1.9636e-04,\n        9.5804e-05, 9.9301e-05, 1.9339e-04, 1.2896e-04, 2.3814e-04, 1.3621e-04,\n        1.4523e-04, 2.1126e-04], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([ 5.9179e-03,  2.4288e-03, -3.1620e-03,  5.0009e-04, -6.5902e-03,\n         2.2251e-03, -1.0909e-03,  8.8967e-04,  2.2133e-03,  3.3763e-03,\n        -3.7177e-03,  1.5845e-02, -2.0845e-03, -3.6139e-03, -2.7652e-03,\n         2.0113e-03, -9.7474e-03, -1.9262e-03,  4.5385e-03,  4.0117e-03,\n         4.6705e-03, -7.5213e-04,  3.7789e-03,  2.8298e-03, -3.2070e-03,\n         1.3596e-05,  3.3771e-04, -4.7682e-03, -6.4656e-04, -2.5916e-03,\n         2.6233e-04, -6.7067e-04,  3.5948e-04,  2.8759e-03, -1.5189e-03,\n         2.0348e-03, -1.9037e-03,  1.1980e-02, -1.9229e-03, -5.8410e-03,\n        -1.0025e-02, -8.4091e-04, -1.4543e-03, -2.5279e-03, -7.3695e-04,\n         3.7865e-03,  4.7156e-03,  1.4225e-03,  1.1423e-03,  1.5592e-03,\n        -8.1527e-06,  1.8039e-03,  5.4561e-04,  5.8288e-04,  4.2004e-03,\n         5.1436e-03,  1.2627e-03,  3.5551e-03, -2.5217e-03,  1.5950e-03,\n         2.1115e-03, -1.0162e-03, -1.7712e-03,  2.2889e-02,  1.4403e-05,\n        -3.5693e-03,  7.3837e-03,  1.5567e-03, -4.1638e-03, -3.7842e-03,\n         1.1757e-02,  5.2155e-03,  1.6883e-04, -5.6254e-04,  1.4957e-03,\n        -3.3904e-03,  3.5483e-03,  7.6258e-04,  2.1427e-03, -3.0355e-03,\n         9.3973e-03, -3.8613e-03, -1.6718e-02, -3.4134e-03,  1.0281e-02,\n        -8.5164e-03, -4.5924e-04, -1.9213e-03,  5.6052e-45, -1.2091e-03,\n         2.8456e-03,  2.5367e-03,  2.0716e-03, -7.0803e-03, -1.4098e-03,\n         7.6939e-03,  1.5692e-03,  7.1461e-03,  5.3430e-03, -1.2514e-04,\n        -7.2770e-03,  4.2862e-03,  1.8691e-03,  7.1632e-03,  2.0415e-03,\n        -4.1226e-04, -4.7169e-04, -4.7672e-04,  1.4480e-03, -7.3167e-03,\n         2.9749e-03,  4.5858e-03, -5.6974e-03,  1.7546e-03,  7.0348e-03,\n         1.0602e-03, -4.3732e-03, -3.7841e-03,  3.3265e-03,  2.1837e-03,\n        -1.3303e-03, -4.1551e-03,  1.1376e-03,  6.2110e-03,  7.3238e-03,\n         8.6742e-04, -2.5176e-03, -6.2455e-03,  1.5207e-04, -7.5959e-04,\n        -3.7181e-04, -7.3135e-03,  4.5273e-03, -1.1502e-03, -2.5702e-03,\n        -2.8237e-03, -1.5067e-03, -2.6923e-03,  1.0748e-03,  3.6037e-03,\n        -1.5874e-03,  5.8147e-03,  1.9918e-03, -3.4507e-03, -2.8908e-03,\n         2.4687e-03,  6.2987e-03,  1.6668e-03,  3.5637e-03,  3.5042e-03,\n        -1.8080e-06,  4.1910e-03, -8.0065e-03, -2.5221e-03, -2.5295e-03,\n         3.8011e-03,  4.3212e-03,  5.5294e-03,  1.5811e-03, -1.9477e-03,\n        -1.0212e-02, -1.8673e-03, -4.1675e-03,  1.4097e-03, -2.8962e-03,\n        -1.0165e-03,  4.9469e-04, -2.2209e-03, -5.6626e-04, -1.5930e-03,\n         1.5485e-04,  2.3632e-03,  1.1629e-03, -1.6872e-03, -3.1925e-03,\n         4.9390e-03, -5.4674e-04,  1.7492e-03,  2.6695e-03,  1.3416e-03,\n         3.2587e-03,  3.1850e-03, -1.9013e-02,  8.2295e-03, -9.7193e-04,\n         2.1936e-04,  3.5064e-04, -4.7841e-03, -6.5066e-04,  6.1262e-05,\n         8.8369e-04, -1.9847e-02,  4.5429e-03,  4.4810e-03,  2.6942e-04,\n        -1.9815e-03, -3.6864e-03,  1.2684e-03,  4.0967e-04,  1.1086e-04,\n         1.2825e-05, -1.1649e-02, -4.6809e-04,  3.1181e-03,  3.6983e-03,\n         1.1450e-03,  1.0211e-03,  1.4666e-03, -1.7599e-03,  1.5627e-03,\n         2.3803e-03,  2.5947e-03, -4.0402e-03, -7.3903e-03, -1.4962e-03,\n         4.8489e-04, -7.3299e-03,  2.4453e-03, -1.4570e-03, -2.7618e-04,\n        -3.1651e-03, -5.2647e-04,  2.7606e-04,  6.9595e-04,  4.4364e-03,\n         2.5898e-04, -3.1478e-03, -5.5785e-03,  2.3480e-03, -7.1072e-04,\n        -1.0468e-03, -3.1532e-03,  2.2517e-03,  8.6322e-03, -1.1294e-02,\n         7.2642e-03,  4.9154e-03,  3.8961e-03,  1.5505e-05,  1.6201e-03,\n        -3.5197e-03,  3.0113e-03, -1.4033e-03,  1.2875e-03, -7.7739e-04,\n        -5.0398e-03,  1.5022e-03, -5.8564e-03,  4.1609e-03,  1.9251e-03,\n         1.7825e-03, -4.4920e-03,  3.2338e-03,  2.1431e-04,  3.6387e-03,\n        -3.2419e-04,  1.7206e-03,  5.7629e-03,  1.1179e-03,  3.2301e-03,\n        -6.6040e-03, -1.3769e-03,  3.1471e-03, -2.0679e-04, -6.3024e-04,\n        -1.4908e-02, -6.1845e-03,  1.7791e-03, -2.5586e-03, -1.3206e-03,\n        -8.2482e-04,  1.8436e-03, -3.3002e-03, -1.3945e-03, -1.9181e-03,\n         5.1308e-03,  1.9556e-03,  1.0338e-03,  2.6772e-04,  8.5433e-03,\n         1.8289e-03, -4.1966e-03,  4.0716e-03, -5.3167e-04,  2.6223e-03,\n         5.9214e-03, -2.4009e-03,  4.3300e-03,  3.9352e-03, -3.8254e-03,\n         4.4148e-03, -7.5865e-03, -1.2896e-03, -2.1571e-04,  3.7853e-03,\n         2.7859e-03, -5.3227e-03, -4.2058e-03,  1.6318e-03, -1.1382e-03,\n         1.8690e-03, -1.3206e-02,  2.7102e-03, -8.0671e-04,  3.8812e-04,\n        -2.3706e-03,  1.1240e-03, -3.0376e-03,  1.1678e-03, -4.3038e-03,\n        -1.2439e-02,  1.9518e-03, -2.5336e-03, -2.5811e-03, -2.9344e-03,\n        -2.4017e-03,  3.0653e-03, -5.5375e-03, -1.3871e-03,  1.4715e-03,\n        -1.4199e-03, -1.1652e-03, -2.5365e-03, -2.7541e-03,  6.8057e-03,\n        -6.7668e-03, -1.8281e-03, -4.9102e-03, -1.0285e-03, -5.9168e-04,\n        -2.3601e-03,  1.4646e-03,  2.9193e-03, -3.1496e-03,  7.2395e-04,\n         1.7736e-03,  4.3440e-03,  1.3586e-03, -4.1171e-04,  4.1991e-04,\n        -5.1535e-04,  1.1736e-04,  6.2646e-03,  2.2715e-03,  2.6446e-03,\n        -5.2332e-04,  3.1694e-03, -2.3462e-03, -7.8308e-03, -1.6162e-03,\n         5.5491e-03, -6.0800e-03,  4.8556e-03,  1.9386e-03,  6.7793e-04,\n         2.3286e-03, -5.3296e-03, -1.7029e-03, -1.3067e-03,  9.9814e-04,\n        -1.8645e-03,  2.1908e-04,  9.6970e-04, -6.5636e-04, -1.2228e-03,\n         1.8147e-04, -7.1861e-04, -1.3943e-02, -4.1772e-04, -3.5339e-03,\n        -1.2662e-03,  5.8686e-03, -1.3779e-03, -2.1644e-03,  3.6517e-03,\n         2.1144e-03,  2.9590e-03,  9.8132e-03, -3.6253e-03,  3.1243e-03,\n         1.5647e-03,  2.7748e-04, -5.5938e-03,  1.5962e-03, -1.0462e-03,\n         5.2084e-04,  5.0388e-04, -1.2982e-03,  6.0874e-03,  8.0412e-04,\n         5.2148e-03,  7.1381e-03, -8.1751e-03, -3.9765e-03, -1.4969e-03,\n         6.9788e-03, -3.6847e-03,  9.2441e-03, -6.6328e-04,  3.5143e-04,\n         2.1182e-03,  3.7152e-03,  4.0868e-03,  2.9959e-03, -2.1054e-03,\n        -9.2702e-03,  4.7768e-03, -4.6514e-03, -4.9798e-03, -1.0931e-03,\n        -1.8984e-03,  1.9416e-03,  4.8931e-04, -2.9233e-03,  3.3145e-03,\n        -1.6660e-03, -3.4924e-04, -2.7322e-03,  6.1902e-03,  2.1998e-04,\n         1.6953e-03, -3.4511e-03,  2.0148e-03,  5.1786e-03, -6.4165e-04,\n         2.1296e-03, -1.2231e-03,  8.4293e-03,  1.4846e-03,  3.4701e-03,\n        -7.1356e-03, -1.7724e-03, -8.6714e-04,  2.9067e-03, -2.0829e-03,\n         2.5659e-05, -3.8062e-03,  1.5016e-03,  1.0422e-03,  1.9606e-03,\n        -8.8975e-03, -1.3261e-04,  4.0590e-03, -2.1499e-03,  1.6862e-03,\n         1.7875e-03, -8.2716e-03,  1.0825e-03,  1.7071e-03,  4.5010e-03,\n         1.6516e-03,  1.9078e-03, -2.1212e-03,  1.0434e-03,  6.1096e-03,\n         1.6403e-04, -5.2427e-03, -1.3559e-03, -4.6814e-03, -2.1426e-03,\n         6.5757e-04,  8.3465e-04,  2.6719e-03, -1.1288e-02,  3.7871e-04,\n        -6.7030e-03,  1.6656e-04,  6.0154e-04, -1.3943e-03, -2.1548e-03,\n         2.8857e-03, -3.1166e-03, -5.5567e-03,  2.7103e-04, -1.2102e-02,\n        -5.9273e-03, -2.8968e-03, -4.5755e-03,  1.0069e-03, -5.6810e-04,\n         2.3190e-03,  1.4321e-03, -3.0973e-03, -2.2337e-04,  1.8840e-03,\n        -1.8492e-03,  7.6737e-03, -4.7839e-04, -3.2643e-03,  4.3095e-03,\n         1.2256e-03,  2.0274e-03,  1.8113e-03, -5.7962e-03, -1.3525e-03,\n        -1.0404e-02, -5.1932e-04,  4.2320e-03, -1.3714e-04,  3.9255e-03,\n        -1.3970e-04, -4.5891e-04,  4.0743e-03, -2.3996e-03, -4.1858e-04,\n        -2.0692e-03,  8.3024e-03,  3.6214e-03, -6.1032e-03,  2.1308e-03,\n        -1.6616e-03,  9.3752e-03], device='cuda:0')",
+        "exp_avg_sq": "tensor([7.1175e-05, 1.4666e-04, 1.2013e-04, 1.1989e-04, 1.3125e-04, 9.5474e-05,\n        5.5966e-05, 9.5907e-05, 1.5575e-04, 1.3895e-04, 1.3559e-04, 1.1171e-03,\n        1.4501e-04, 1.5508e-04, 2.1205e-04, 7.2237e-05, 2.1174e-04, 2.2622e-04,\n        2.2056e-04, 1.1653e-04, 1.4728e-04, 8.5511e-05, 1.0299e-04, 6.7795e-05,\n        2.0711e-04, 8.8416e-05, 1.6652e-04, 1.3755e-04, 2.8306e-04, 2.6655e-04,\n        1.0102e-04, 1.5262e-04, 7.5848e-05, 6.6357e-05, 1.8992e-04, 9.9344e-05,\n        1.1777e-04, 1.9557e-04, 1.4310e-04, 9.0165e-05, 1.1001e-04, 1.1800e-04,\n        1.2269e-04, 1.4745e-04, 7.8662e-05, 1.1200e-04, 1.3440e-04, 1.7420e-04,\n        2.9178e-04, 8.0590e-05, 6.6216e-04, 2.9953e-04, 1.3417e-04, 8.2108e-05,\n        1.5118e-04, 1.1411e-04, 1.3849e-04, 1.8482e-04, 2.0128e-04, 2.1895e-04,\n        2.1034e-04, 1.2737e-04, 2.1742e-04, 7.0049e-03, 1.2139e-04, 1.3009e-04,\n        1.7836e-04, 1.4483e-04, 5.6493e-05, 2.1242e-04, 9.5870e-05, 1.9338e-04,\n        1.5128e-04, 1.0852e-04, 1.0317e-04, 2.8749e-04, 9.3913e-05, 3.8385e-04,\n        2.0008e-04, 1.5475e-04, 2.8746e-04, 5.1517e-04, 5.8982e-04, 1.2836e-04,\n        1.9666e-04, 8.2556e-05, 1.0065e-04, 1.7549e-04, 4.2788e-09, 2.3508e-04,\n        1.0242e-04, 2.1860e-04, 1.3710e-04, 1.5596e-04, 3.2840e-04, 1.4293e-04,\n        8.9726e-05, 4.8278e-04, 1.0964e-04, 2.5211e-04, 1.7257e-04, 9.3426e-05,\n        1.0423e-04, 3.5303e-04, 1.6504e-04, 7.8838e-05, 8.4499e-05, 9.1568e-05,\n        1.0345e-04, 1.8069e-04, 2.0385e-04, 1.1273e-04, 1.3993e-04, 6.7499e-05,\n        1.2605e-04, 1.1379e-04, 1.1638e-04, 1.2924e-04, 1.0915e-04, 5.0988e-05,\n        1.9766e-04, 9.7728e-05, 1.7015e-04, 2.0158e-04, 1.4570e-04, 8.8572e-05,\n        7.3807e-05, 7.1639e-05, 9.3599e-05, 2.1319e-04, 1.1915e-04, 1.8966e-04,\n        1.6046e-04, 2.4895e-04, 3.3983e-04, 1.9387e-04, 1.1589e-04, 1.2784e-04,\n        9.0809e-05, 2.2186e-04, 1.4346e-04, 2.0086e-04, 1.2570e-04, 1.6064e-04,\n        1.6037e-04, 1.4611e-04, 2.5658e-04, 1.3963e-04, 1.0155e-04, 1.5909e-04,\n        7.9474e-05, 1.7451e-04, 1.0883e-04, 1.5226e-04, 9.0137e-05, 1.4324e-04,\n        1.2482e-04, 5.6658e-05, 1.3776e-04, 1.0966e-04, 1.0330e-04, 1.0768e-04,\n        1.9002e-04, 1.0043e-04, 1.1680e-04, 1.7373e-04, 7.2426e-05, 2.2572e-04,\n        1.5279e-04, 2.1080e-04, 1.2814e-04, 1.9749e-04, 1.4092e-04, 1.2123e-04,\n        3.7377e-04, 1.2019e-04, 1.5099e-04, 2.5081e-04, 1.9434e-04, 2.0892e-04,\n        1.1456e-04, 1.5322e-04, 2.2800e-04, 4.3026e-04, 1.9856e-04, 1.0693e-04,\n        2.2683e-04, 3.8972e-04, 1.4191e-04, 1.8889e-04, 1.1399e-04, 2.3602e-04,\n        2.3678e-04, 7.1652e-05, 1.7209e-04, 7.8757e-05, 1.6825e-04, 1.1253e-04,\n        1.1208e-04, 7.7576e-05, 1.5685e-04, 2.7045e-04, 2.2336e-04, 1.4308e-04,\n        1.4151e-04, 1.9520e-04, 1.1178e-04, 1.1497e-04, 6.9835e-05, 9.4069e-05,\n        1.1073e-04, 2.8151e-04, 2.0439e-04, 1.1097e-04, 1.4920e-04, 1.4431e-04,\n        8.9549e-05, 2.1006e-04, 3.0436e-04, 1.6202e-04, 2.0814e-04, 1.3489e-04,\n        2.0014e-04, 1.6465e-04, 1.8940e-04, 5.5719e-05, 2.2733e-04, 9.9278e-05,\n        1.4673e-04, 1.0007e-04, 4.8407e-05, 1.1687e-04, 1.8682e-04, 2.3849e-04,\n        2.8017e-04, 1.3796e-04, 2.7824e-04, 2.2457e-04, 6.2522e-05, 8.8890e-05,\n        1.5642e-04, 7.5219e-05, 8.8979e-05, 1.9531e-04, 1.0501e-04, 1.1453e-04,\n        7.6857e-05, 1.3278e-04, 3.7375e-04, 1.3996e-04, 2.0352e-04, 1.7106e-04,\n        1.6363e-04, 1.1508e-04, 1.1408e-04, 1.5128e-04, 7.2099e-05, 2.3732e-04,\n        6.0973e-05, 1.1759e-04, 1.3243e-04, 2.3874e-04, 1.3828e-04, 9.0596e-05,\n        1.6268e-04, 2.4502e-04, 1.5468e-04, 7.7388e-05, 9.5793e-05, 1.8899e-04,\n        1.1937e-04, 9.3243e-05, 1.7991e-04, 1.0463e-04, 1.2113e-04, 4.6656e-04,\n        1.7038e-04, 1.4686e-04, 1.9492e-04, 2.2054e-04, 1.5912e-04, 1.5327e-04,\n        3.6605e-04, 1.2268e-04, 1.2557e-04, 1.5061e-04, 1.0913e-04, 2.8510e-04,\n        2.3443e-04, 2.5786e-04, 7.5165e-05, 1.8867e-04, 1.0341e-04, 7.2293e-05,\n        1.5550e-04, 1.7008e-04, 2.7282e-04, 3.4179e-04, 5.2470e-05, 2.9169e-04,\n        1.2613e-04, 7.0636e-04, 8.7767e-05, 7.6629e-05, 1.1408e-04, 2.3363e-04,\n        2.2286e-04, 1.9920e-04, 6.7311e-05, 1.2537e-04, 2.5613e-04, 1.1757e-04,\n        8.6190e-05, 1.3787e-04, 7.0569e-05, 8.5019e-05, 2.5395e-04, 2.1999e-04,\n        1.3997e-04, 1.2824e-04, 2.3120e-04, 2.6346e-04, 1.6609e-04, 1.5159e-04,\n        2.8784e-04, 9.5537e-05, 1.2665e-04, 1.1337e-04, 1.7194e-04, 1.9314e-04,\n        1.7271e-04, 2.1445e-04, 1.5554e-04, 1.1803e-04, 9.4835e-05, 1.2489e-04,\n        9.7879e-05, 1.6447e-04, 5.9460e-05, 1.8565e-04, 1.9044e-04, 1.4530e-04,\n        1.3499e-04, 1.8377e-04, 9.1965e-05, 2.2272e-04, 2.4216e-04, 1.4595e-04,\n        7.3163e-05, 3.2908e-04, 1.0088e-04, 2.3814e-04, 3.2496e-04, 7.9750e-05,\n        1.2595e-04, 2.6875e-04, 2.4396e-04, 1.3370e-04, 2.2907e-04, 8.1449e-05,\n        1.5257e-04, 7.4695e-05, 5.9321e-05, 1.7478e-04, 1.0921e-04, 4.0271e-05,\n        9.1579e-05, 3.6209e-04, 2.1234e-04, 1.2756e-04, 1.6142e-04, 1.3557e-04,\n        1.1898e-04, 1.4819e-04, 1.3604e-04, 3.6208e-04, 6.9492e-05, 3.9082e-04,\n        1.0651e-04, 2.2054e-04, 1.6078e-04, 1.2716e-04, 2.1047e-04, 2.6919e-04,\n        1.9055e-04, 3.3430e-04, 1.0648e-04, 1.3429e-04, 1.8872e-04, 1.7697e-04,\n        1.4761e-04, 4.4019e-04, 9.1675e-05, 1.8625e-04, 1.6765e-04, 1.3441e-04,\n        8.4081e-05, 2.5302e-04, 1.7665e-04, 5.8364e-05, 1.4961e-04, 1.2196e-04,\n        1.7554e-04, 9.9895e-05, 8.6891e-05, 2.9479e-04, 1.0218e-04, 2.0177e-04,\n        9.6523e-05, 1.6396e-04, 1.8132e-04, 1.7522e-04, 9.2247e-05, 1.2239e-04,\n        1.1903e-04, 1.4660e-04, 1.1903e-04, 1.7398e-04, 4.2109e-04, 1.7038e-04,\n        1.4555e-04, 2.9065e-04, 1.5793e-04, 1.7824e-04, 1.4013e-04, 1.4241e-04,\n        9.7378e-05, 1.8195e-04, 1.3112e-04, 8.1108e-05, 8.9771e-05, 1.3310e-04,\n        1.1886e-04, 6.8161e-04, 1.2090e-04, 1.9303e-04, 2.9866e-04, 1.3166e-04,\n        1.5539e-04, 2.4175e-04, 4.0524e-04, 9.9937e-05, 1.0646e-04, 2.1116e-04,\n        1.4154e-04, 1.9992e-04, 1.8720e-04, 1.7202e-04, 1.4326e-04, 2.6748e-04,\n        3.2066e-04, 1.2660e-04, 8.3293e-05, 2.7975e-04, 1.5592e-04, 2.4764e-04,\n        1.1224e-04, 2.4504e-04, 1.7405e-04, 1.4489e-04, 1.7651e-04, 3.3916e-04,\n        2.2134e-04, 2.0007e-04, 1.8501e-04, 4.1757e-04, 1.6057e-04, 1.7545e-04,\n        1.6032e-04, 1.5428e-04, 2.3689e-04, 1.1315e-04, 6.6128e-05, 1.4059e-04,\n        1.7540e-04, 1.2061e-04, 2.4748e-04, 1.7512e-04, 5.8312e-05, 1.1535e-04,\n        8.3178e-05, 2.1024e-04, 2.0771e-04, 6.7370e-05, 7.4603e-05, 9.3156e-05,\n        1.2705e-04, 1.1795e-04, 1.2496e-04, 8.2481e-05, 1.8947e-04, 1.6479e-04,\n        1.2039e-04, 2.1047e-04, 1.1586e-04, 1.5753e-04, 1.2669e-04, 1.0907e-04,\n        7.0165e-05, 3.7207e-05, 2.1764e-04, 7.4360e-05, 1.2548e-04, 2.0052e-04,\n        8.7653e-05, 9.4995e-05, 1.9266e-04, 1.2327e-04, 2.4601e-04, 1.3512e-04,\n        1.4669e-04, 2.0777e-04], device='cuda:0')"
       },
       "3": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([-1.0853e-03, -9.8263e-03,  1.3180e-03, -6.4797e-04, -4.7183e-03,\n        -6.7327e-04, -3.0160e-05,  2.3715e-03,  3.2611e-03, -3.1801e-03,\n        -2.4556e-03, -4.1822e-04,  3.3329e-03, -1.8064e-03,  4.1467e-04,\n        -1.2677e-03,  4.9186e-03,  6.2149e-04,  4.7019e-03, -1.1823e-03,\n        -6.8856e-04, -1.1852e-03,  7.7870e-04,  8.3500e-04, -4.5657e-03,\n        -6.1391e-04,  1.3233e-03, -9.7638e-04,  3.5372e-04, -6.3919e-03,\n        -8.4232e-04, -7.1042e-03,  7.8620e-04, -5.3699e-04, -1.8762e-03,\n         8.7332e-05,  1.9962e-04,  1.8765e-03, -5.7674e-04, -1.8008e-03,\n         1.1337e-03, -2.5889e-04, -1.8274e-03, -1.9650e-03, -2.5319e-03,\n        -1.7920e-03,  4.6591e-04,  3.9217e-03,  1.5644e-03,  4.3023e-04,\n        -1.2159e-03, -1.5490e-03, -5.6399e-04, -1.0458e-03, -8.3199e-04,\n         3.6179e-04,  2.9003e-03, -3.0994e-03,  3.0558e-03,  2.8557e-03,\n         1.4931e-03,  1.2768e-03, -3.6877e-03, -8.0495e-03,  2.4656e-03,\n         2.1318e-03, -2.5042e-03, -2.0798e-03, -5.1605e-04,  9.3056e-04,\n         1.4211e-03, -1.2400e-03, -4.5766e-03,  6.2149e-05,  1.0827e-03,\n        -6.0184e-04, -1.4343e-03,  9.3775e-04,  1.7490e-03,  5.6686e-05,\n        -3.9882e-03, -3.0967e-03, -3.1115e-03,  3.6919e-04, -1.1828e-03,\n        -4.1629e-04, -3.6189e-03,  2.2371e-03,  5.6052e-45,  1.8652e-03,\n        -4.3921e-03, -3.3353e-03,  4.0175e-04, -2.0453e-03,  8.4775e-03,\n        -1.4714e-03, -1.3530e-03,  2.4065e-03, -4.4827e-04, -3.9348e-03,\n         5.1154e-04, -1.4036e-03, -1.2052e-03, -2.3825e-03, -3.6663e-03,\n         2.3216e-03, -2.0341e-03,  2.1561e-03, -1.7072e-03, -3.0870e-04,\n         3.7181e-04, -1.1656e-03, -4.4529e-04,  4.7219e-04, -5.3375e-04,\n         3.6017e-03, -2.6905e-04,  5.7839e-05,  1.6978e-04, -2.3465e-04,\n         2.3591e-03,  2.6312e-03,  9.4818e-04,  9.2773e-04, -4.3627e-03,\n        -5.4488e-03, -4.2867e-04, -7.7594e-04,  2.9132e-03, -1.2536e-04,\n        -6.9944e-04, -2.0893e-03, -1.8206e-03, -6.0117e-04,  2.5988e-04,\n        -6.3356e-04,  9.8817e-04,  1.4934e-03, -5.9880e-04, -4.9690e-03,\n         1.8512e-03, -2.0788e-03, -1.6956e-03, -2.9854e-03, -1.7401e-03,\n        -2.0141e-03,  3.9103e-04, -6.9497e-04, -8.8307e-05, -3.1481e-03,\n        -4.4267e-03, -7.6877e-04, -8.4378e-04,  1.1194e-03, -1.1963e-03,\n        -3.4301e-03,  1.2954e-03,  1.4214e-03,  3.7362e-03,  1.4824e-03,\n         3.6561e-04, -1.4134e-03, -6.7820e-04, -1.4220e-04,  5.0719e-05,\n         1.6444e-03,  7.7188e-04,  1.1119e-03,  5.6631e-04,  1.3627e-04,\n         8.4947e-04,  2.3133e-03, -1.6362e-03,  5.2457e-04, -8.0766e-05,\n        -6.8920e-04,  3.6448e-03,  1.4235e-03,  3.9349e-03, -4.1745e-03,\n         1.9179e-03,  8.9470e-04, -1.1892e-03,  1.4328e-03,  1.2704e-03,\n        -2.0201e-03,  6.9273e-04, -1.8164e-03, -2.7076e-03,  1.7869e-03,\n         1.5623e-03, -8.6988e-04,  5.2896e-03, -7.1698e-04,  2.4343e-04,\n         1.6051e-03,  1.9848e-03, -8.0977e-04, -6.9233e-04,  1.7901e-03,\n         9.5112e-04, -6.0550e-03,  1.2146e-04, -6.5151e-03, -3.1950e-04,\n        -1.7611e-04, -1.1667e-03, -7.5443e-04,  8.9804e-04, -2.6099e-03,\n         4.2195e-04,  1.3118e-03,  3.1940e-03,  1.0206e-03,  2.2467e-03,\n         3.1839e-03,  1.9295e-03, -1.0884e-04,  6.7329e-04,  4.1733e-03,\n         1.7125e-03,  3.7820e-03, -3.0822e-04, -1.5643e-03,  1.7603e-03,\n         2.2982e-03,  1.3975e-04,  7.4328e-04,  1.7721e-03, -7.3573e-04,\n        -1.3709e-05,  3.2038e-04,  1.1714e-03,  1.8234e-03, -1.1164e-03,\n         3.1238e-04,  2.6781e-03, -1.2185e-03,  3.6563e-04, -1.4720e-03,\n         4.8017e-03, -2.0991e-03,  1.3631e-03, -1.2494e-03, -1.6484e-03,\n        -1.5662e-03, -6.2174e-04,  3.2116e-03, -2.6574e-04,  6.5042e-04,\n        -3.9142e-03, -1.1072e-03, -2.1614e-03, -2.4905e-04, -1.0648e-04,\n         1.3890e-03,  2.7925e-04, -4.7628e-03, -1.3568e-03,  1.0523e-03,\n         3.2957e-05, -2.5272e-04,  1.8212e-03,  6.1748e-04,  5.2582e-04,\n         1.5793e-03,  8.6274e-04, -2.5888e-03, -1.7527e-04, -1.2068e-03,\n        -2.3251e-03,  1.3547e-03,  1.2032e-03,  2.8845e-05,  1.0216e-03,\n         3.6469e-03,  2.1972e-03, -1.1690e-03, -3.5749e-03, -9.7676e-04,\n         4.6588e-04,  1.9913e-03,  1.7301e-03,  3.1914e-05,  3.6072e-03,\n         3.8759e-03,  1.1001e-03,  3.4442e-05,  1.8187e-03,  2.0290e-04,\n         1.8916e-03,  1.2604e-03,  3.5289e-04, -4.5221e-04, -7.3355e-03,\n        -2.2403e-03, -2.2409e-03,  1.7212e-03,  1.9050e-03,  3.8709e-03,\n         1.4710e-04,  2.9055e-03,  1.8332e-03, -4.8738e-03,  1.0475e-03,\n        -3.4267e-03, -5.6241e-03, -4.0356e-04, -1.2768e-03, -2.8577e-04,\n         1.0235e-03, -2.4586e-04, -5.6762e-04, -1.7885e-03,  1.4649e-03,\n         1.6406e-04, -6.1456e-04,  1.1210e-03, -5.7546e-03, -4.4731e-04,\n        -7.8874e-04,  1.5838e-03, -5.0258e-04, -3.0063e-03,  6.8658e-04,\n        -1.1009e-04, -3.2345e-03, -3.3587e-03, -1.3573e-05,  1.5878e-03,\n         3.2488e-03,  3.8578e-03, -2.4640e-04, -1.2933e-03,  1.9635e-03,\n        -2.8860e-03, -2.6748e-03, -7.9620e-04, -2.3501e-03, -2.8541e-04,\n        -1.0913e-03, -8.9856e-04,  3.0051e-03,  6.4363e-04,  4.4462e-04,\n        -6.0399e-03,  1.8520e-03, -6.0454e-04, -4.0802e-04,  1.8385e-03,\n         2.2400e-03, -1.9480e-03,  1.3909e-03, -2.7998e-03, -3.6891e-03,\n        -1.8812e-03, -2.4458e-03, -3.3900e-05, -2.6193e-03,  1.0586e-03,\n         6.2476e-03,  9.4603e-04, -7.0110e-04, -2.1305e-03,  7.3628e-05,\n        -6.9113e-04, -1.9844e-03,  8.0640e-03, -6.7008e-03,  4.2359e-04,\n        -2.9381e-03,  1.2464e-03, -2.0340e-04, -1.8108e-04, -2.6328e-03,\n        -1.3416e-04,  5.7107e-04,  4.1640e-03,  3.3799e-03,  2.7257e-03,\n        -7.7970e-04, -2.0331e-03,  1.5036e-05,  2.3361e-03,  8.1594e-04,\n        -1.6032e-03,  2.9209e-04, -2.1746e-03,  9.4527e-04, -5.9377e-03,\n         1.4791e-03,  2.2397e-03,  2.6918e-03,  1.0580e-03, -2.2186e-03,\n        -3.3350e-03, -1.6416e-04, -2.2035e-03,  1.3752e-03, -4.1342e-03,\n        -1.4359e-03,  2.0829e-03,  1.7854e-03, -4.2486e-04, -1.1840e-03,\n         2.3372e-03, -3.7411e-04,  4.8416e-04,  2.3816e-03,  3.1573e-04,\n         1.2482e-03,  2.9099e-04,  2.8379e-03,  1.9299e-03,  1.1633e-03,\n        -3.1250e-04,  1.6705e-03,  8.9216e-04,  6.3172e-03,  3.8281e-03,\n         4.0878e-04,  4.5281e-03, -1.3737e-03, -2.3228e-03,  1.1338e-03,\n         1.3025e-03,  8.4158e-04,  1.0279e-03,  3.5120e-03, -1.2751e-05,\n        -4.3223e-03,  3.1727e-04, -1.4560e-03,  3.5569e-03,  2.4044e-03,\n        -1.6422e-04, -9.6495e-03, -4.8815e-03,  1.2022e-03, -2.4302e-03,\n         9.4156e-04, -9.7439e-04,  2.1443e-03,  1.4621e-03, -1.6774e-04,\n        -1.3023e-03,  5.4488e-05,  9.9117e-04, -3.9377e-04,  1.9675e-03,\n         2.4495e-03, -3.7126e-04, -2.2104e-03,  1.9258e-03,  1.7710e-03,\n        -2.1678e-03,  1.9849e-03,  1.5981e-03, -2.4690e-03,  1.8752e-03,\n        -7.8588e-04, -8.2233e-04, -1.4671e-03, -2.1319e-03, -3.1580e-04,\n         9.2874e-04,  5.8542e-03,  2.3137e-03,  1.6218e-03,  6.7673e-03,\n        -4.4859e-03, -2.3119e-03, -1.1446e-03, -8.5892e-04,  2.2350e-03,\n         1.1087e-03,  1.0696e-03,  8.1584e-04, -1.2218e-04, -1.8899e-04,\n         2.4442e-03,  9.0573e-03,  2.7737e-03, -1.0695e-03,  2.5903e-03,\n        -1.8094e-03, -3.1225e-03, -1.5753e-03,  8.9554e-04,  1.1178e-03,\n         1.0740e-03,  6.5198e-04,  1.5208e-03,  2.4760e-03,  1.7397e-03,\n         1.3183e-03,  1.7894e-03,  3.7913e-04,  1.6786e-03,  5.9341e-05,\n        -3.0455e-03, -1.4465e-03,  2.4962e-03,  1.8052e-03, -1.1286e-04,\n        -2.5528e-03, -2.0758e-03,  5.7943e-04, -1.5952e-03,  1.2261e-03,\n         1.2478e-04, -1.5115e-03], device='cuda:0')",
-        "exp_avg_sq": "tensor([3.6564e-05, 7.5405e-05, 5.8471e-05, 4.5447e-05, 4.2688e-05, 3.9452e-05,\n        2.7969e-05, 3.2697e-05, 5.0900e-05, 4.1646e-05, 4.7861e-05, 4.8659e-05,\n        6.8098e-05, 7.9423e-05, 5.4315e-05, 3.9878e-05, 6.5369e-05, 5.5223e-05,\n        6.4389e-05, 5.3265e-05, 5.0648e-05, 4.1356e-05, 5.5710e-05, 3.1117e-05,\n        9.4576e-05, 3.6929e-05, 4.8472e-05, 7.3248e-05, 6.3705e-05, 5.7682e-05,\n        5.3504e-05, 6.6247e-05, 3.0762e-05, 3.3515e-05, 8.7140e-05, 3.0366e-05,\n        7.0178e-05, 7.4690e-05, 3.0766e-05, 5.3167e-05, 3.9646e-05, 4.7899e-05,\n        4.2239e-05, 4.8917e-05, 3.5626e-05, 3.9879e-05, 4.6460e-05, 4.4233e-05,\n        5.8573e-05, 3.1586e-05, 5.3416e-05, 7.0513e-05, 4.0645e-05, 5.7062e-05,\n        4.3344e-05, 4.1265e-05, 4.0468e-05, 4.9004e-05, 5.3290e-05, 6.4904e-05,\n        6.5762e-05, 4.4054e-05, 6.0575e-05, 1.5367e-04, 5.0398e-05, 4.6063e-05,\n        6.1859e-05, 5.2021e-05, 2.8603e-05, 6.1602e-05, 5.2366e-05, 6.7410e-05,\n        5.1351e-05, 4.2244e-05, 3.1786e-05, 9.5083e-05, 4.4472e-05, 8.7723e-05,\n        4.8082e-05, 5.6849e-05, 6.5025e-05, 7.3297e-05, 8.3224e-05, 5.8023e-05,\n        6.0802e-05, 4.6987e-05, 6.0213e-05, 8.3307e-05, 2.2474e-10, 7.3602e-05,\n        5.2241e-05, 1.0037e-04, 4.4129e-05, 4.6414e-05, 9.3487e-05, 4.5099e-05,\n        4.0907e-05, 5.1609e-05, 4.8744e-05, 6.1484e-05, 4.2488e-05, 4.1982e-05,\n        3.6180e-05, 8.2992e-05, 5.7640e-05, 4.7112e-05, 4.3073e-05, 4.0830e-05,\n        4.3324e-05, 5.1858e-05, 4.4800e-05, 3.8260e-05, 6.0918e-05, 5.4886e-05,\n        4.2167e-05, 3.9528e-05, 5.8057e-05, 4.9992e-05, 3.9775e-05, 2.8654e-05,\n        6.1397e-05, 4.7852e-05, 5.7892e-05, 5.3299e-05, 3.7136e-05, 4.1363e-05,\n        3.1480e-05, 2.9930e-05, 4.2899e-05, 4.3869e-05, 3.9970e-05, 6.7370e-05,\n        5.4243e-05, 6.6455e-05, 7.5681e-05, 6.8014e-05, 4.2326e-05, 5.7583e-05,\n        3.7718e-05, 5.0999e-05, 5.8303e-05, 4.6931e-05, 5.4304e-05, 6.0096e-05,\n        5.4155e-05, 6.4782e-05, 7.5907e-05, 5.1695e-05, 4.7262e-05, 6.9972e-05,\n        4.0760e-05, 4.3162e-05, 4.4414e-05, 4.2235e-05, 4.6027e-05, 4.9909e-05,\n        4.8694e-05, 2.7738e-05, 6.9443e-05, 4.8918e-05, 4.3388e-05, 3.0855e-05,\n        4.5253e-05, 5.0845e-05, 2.7946e-05, 6.2479e-05, 5.0000e-05, 6.7770e-05,\n        5.0454e-05, 6.1745e-05, 4.3562e-05, 5.7830e-05, 5.8049e-05, 5.8639e-05,\n        8.3398e-05, 5.1883e-05, 5.3628e-05, 6.5211e-05, 5.1086e-05, 8.6472e-05,\n        6.8430e-05, 6.0386e-05, 4.7202e-05, 1.0935e-04, 7.3422e-05, 3.5149e-05,\n        6.6598e-05, 8.7330e-05, 4.8819e-05, 6.2356e-05, 4.8408e-05, 8.7827e-05,\n        6.8855e-05, 3.5736e-05, 6.5443e-05, 3.8172e-05, 4.8884e-05, 3.1957e-05,\n        4.2938e-05, 3.1233e-05, 6.4824e-05, 6.0251e-05, 4.9597e-05, 5.9628e-05,\n        3.7390e-05, 5.2807e-05, 3.9028e-05, 3.8293e-05, 3.4981e-05, 5.6005e-05,\n        3.8040e-05, 6.3048e-05, 7.1086e-05, 4.0791e-05, 7.4870e-05, 4.6079e-05,\n        4.2015e-05, 5.4863e-05, 6.5278e-05, 5.8183e-05, 6.5748e-05, 8.3060e-05,\n        6.7008e-05, 3.9016e-05, 5.7434e-05, 3.7725e-05, 7.0557e-05, 4.2469e-05,\n        4.7866e-05, 4.3565e-05, 2.6769e-05, 4.8149e-05, 5.4846e-05, 6.2357e-05,\n        6.2316e-05, 4.1340e-05, 9.0705e-05, 8.4854e-05, 3.7702e-05, 3.5962e-05,\n        5.6524e-05, 3.4793e-05, 3.3909e-05, 5.0210e-05, 3.3112e-05, 4.8244e-05,\n        3.6219e-05, 5.0557e-05, 6.0089e-05, 3.7722e-05, 5.7076e-05, 4.6353e-05,\n        6.3125e-05, 4.2811e-05, 4.4377e-05, 4.2003e-05, 3.0775e-05, 7.8677e-05,\n        3.8073e-05, 4.8535e-05, 5.1669e-05, 6.6932e-05, 6.4008e-05, 4.2516e-05,\n        4.1909e-05, 7.0325e-05, 4.0027e-05, 3.8751e-05, 4.0046e-05, 6.4787e-05,\n        4.1985e-05, 3.0508e-05, 6.4464e-05, 3.7598e-05, 4.3402e-05, 7.9170e-05,\n        5.6396e-05, 3.8723e-05, 6.6585e-05, 5.5863e-05, 6.0161e-05, 4.4350e-05,\n        8.7441e-05, 6.4187e-05, 4.1854e-05, 6.2886e-05, 4.1353e-05, 7.3304e-05,\n        6.4213e-05, 4.4431e-05, 3.6533e-05, 1.0839e-04, 4.7311e-05, 3.9573e-05,\n        5.0391e-05, 4.3731e-05, 1.0242e-04, 1.0995e-04, 2.3910e-05, 6.3379e-05,\n        3.8653e-05, 7.9845e-05, 3.4554e-05, 3.8530e-05, 5.9207e-05, 8.0649e-05,\n        7.0031e-05, 4.4081e-05, 4.2080e-05, 4.6812e-05, 5.2697e-05, 4.3482e-05,\n        3.6584e-05, 4.4348e-05, 2.7366e-05, 2.9805e-05, 6.3668e-05, 5.6828e-05,\n        4.7297e-05, 4.9803e-05, 5.8514e-05, 8.7063e-05, 7.6416e-05, 7.2209e-05,\n        6.9007e-05, 3.3576e-05, 4.3152e-05, 4.8740e-05, 4.2296e-05, 5.9548e-05,\n        5.3495e-05, 6.9777e-05, 5.9488e-05, 4.7749e-05, 3.0667e-05, 4.8876e-05,\n        5.2055e-05, 5.2724e-05, 4.3486e-05, 5.8208e-05, 5.7690e-05, 5.4881e-05,\n        4.4490e-05, 5.7082e-05, 4.9835e-05, 8.2346e-05, 6.4914e-05, 5.0745e-05,\n        3.3225e-05, 7.1799e-05, 4.0396e-05, 6.8564e-05, 5.6974e-05, 4.7553e-05,\n        5.3553e-05, 6.8612e-05, 5.5354e-05, 5.3846e-05, 7.0704e-05, 3.1546e-05,\n        6.7693e-05, 2.7365e-05, 3.5609e-05, 5.4473e-05, 5.0420e-05, 2.5218e-05,\n        4.1292e-05, 8.6843e-05, 7.3149e-05, 3.8347e-05, 7.9377e-05, 4.7222e-05,\n        4.9269e-05, 3.1135e-05, 6.1338e-05, 1.2889e-04, 3.8226e-05, 7.9283e-05,\n        4.7117e-05, 4.5292e-05, 5.7085e-05, 4.8115e-05, 1.2664e-04, 6.5484e-05,\n        5.8772e-05, 6.9729e-05, 3.2784e-05, 6.9423e-05, 5.6744e-05, 8.6470e-05,\n        6.4042e-05, 6.2655e-05, 3.9613e-05, 4.4865e-05, 5.1234e-05, 5.1574e-05,\n        3.7914e-05, 5.7553e-05, 5.1880e-05, 3.2619e-05, 7.5285e-05, 7.0538e-05,\n        5.1978e-05, 3.0031e-05, 3.4364e-05, 9.0259e-05, 4.4613e-05, 4.2894e-05,\n        4.7071e-05, 4.8162e-05, 5.3888e-05, 5.7253e-05, 4.8146e-05, 4.8578e-05,\n        5.7756e-05, 4.0932e-05, 5.3400e-05, 7.5246e-05, 6.6350e-05, 7.7586e-05,\n        4.7523e-05, 8.2966e-05, 6.0166e-05, 5.6771e-05, 5.6806e-05, 6.0543e-05,\n        4.2586e-05, 5.6575e-05, 4.6760e-05, 5.0156e-05, 3.8497e-05, 4.6196e-05,\n        3.7944e-05, 9.1421e-05, 4.4083e-05, 6.0409e-05, 8.5052e-05, 5.6153e-05,\n        6.1571e-05, 7.0128e-05, 1.2537e-04, 3.7887e-05, 5.0234e-05, 5.5660e-05,\n        4.8116e-05, 5.6766e-05, 5.7520e-05, 4.4245e-05, 5.9545e-05, 5.9759e-05,\n        9.1233e-05, 4.6183e-05, 4.6349e-05, 5.8792e-05, 5.0922e-05, 6.4573e-05,\n        2.9241e-05, 5.4557e-05, 4.9860e-05, 3.9185e-05, 5.9437e-05, 6.5517e-05,\n        5.3694e-05, 5.7463e-05, 7.4153e-05, 8.5458e-05, 6.4916e-05, 7.5933e-05,\n        6.9330e-05, 4.9410e-05, 9.4836e-05, 3.2556e-05, 3.0617e-05, 5.4766e-05,\n        5.7558e-05, 4.0870e-05, 4.1419e-05, 4.7886e-05, 4.1166e-05, 4.7082e-05,\n        4.7489e-05, 7.8927e-05, 6.3693e-05, 3.1705e-05, 3.8194e-05, 5.4497e-05,\n        6.9612e-05, 5.1421e-05, 5.6470e-05, 2.7790e-05, 7.1239e-05, 5.0059e-05,\n        3.9554e-05, 6.0055e-05, 4.4640e-05, 6.6643e-05, 4.0433e-05, 3.6323e-05,\n        3.7176e-05, 1.6234e-05, 7.1945e-05, 4.3831e-05, 4.0362e-05, 5.0363e-05,\n        4.1242e-05, 4.5514e-05, 5.9081e-05, 4.4081e-05, 8.4834e-05, 4.3894e-05,\n        5.3950e-05, 5.0730e-05], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([ 4.1373e-03,  2.0046e-03, -2.0093e-03,  6.0270e-04, -4.2223e-03,\n         1.4737e-03, -8.3424e-04,  9.9720e-04,  1.3730e-03,  2.3100e-03,\n        -1.5413e-03,  3.5469e-03, -9.4580e-04, -2.4875e-03, -1.5314e-03,\n         1.3928e-03, -5.4596e-03, -9.6337e-04,  2.7346e-03,  2.4711e-03,\n         2.5775e-03, -3.8221e-04,  2.5814e-03,  1.7790e-03, -1.9503e-03,\n         1.3163e-04,  1.7431e-04, -2.8731e-03, -2.3889e-04, -1.5291e-03,\n         4.2299e-04, -5.5862e-05,  5.5053e-04,  2.2228e-03, -1.0593e-03,\n         1.3684e-03, -1.3330e-03,  6.9733e-03, -8.4587e-04, -5.1940e-03,\n        -5.5081e-03, -1.6495e-04, -7.0945e-04, -1.2496e-03, -4.8574e-04,\n         2.2032e-03,  2.6144e-03,  1.0518e-03,  4.1722e-04,  1.1655e-03,\n         4.6709e-05,  4.1575e-04,  3.0638e-04,  3.0734e-04,  1.9658e-03,\n         3.8384e-03,  7.7699e-04,  1.5948e-03, -1.2091e-03,  1.6686e-03,\n         1.5183e-03, -4.7683e-04, -4.1525e-04,  3.2692e-03, -2.4660e-04,\n        -2.1823e-03,  4.1902e-03,  1.1494e-03, -2.8962e-03, -1.8233e-03,\n         7.8810e-03,  2.8802e-03, -2.3219e-05, -7.2163e-04,  5.2711e-04,\n        -2.0807e-03,  2.9093e-03, -2.8891e-04,  8.1500e-04, -1.8922e-03,\n         3.2743e-03, -1.3520e-03, -6.1703e-03, -2.3846e-03,  5.1012e-03,\n        -6.7753e-03, -2.0012e-04, -1.3147e-03,  5.6052e-45, -7.7328e-04,\n         1.3411e-03,  1.7005e-03,  1.1283e-03, -3.3912e-03, -8.7438e-04,\n         4.0185e-03,  1.2406e-03,  2.2250e-03,  3.4867e-03, -1.1898e-04,\n        -3.9281e-03,  1.7390e-03,  1.3889e-03,  3.9410e-03,  1.2456e-03,\n        -4.8506e-04, -1.4997e-04,  9.4303e-05,  6.5614e-04, -3.2496e-03,\n         1.9876e-03,  2.3487e-03, -3.6933e-03,  9.7908e-04,  4.1263e-03,\n         8.8741e-04, -3.6422e-03, -2.4905e-03,  2.2794e-03,  1.6183e-03,\n         3.6154e-05, -3.4959e-03,  9.0872e-04,  3.3306e-03,  3.5380e-03,\n        -3.7241e-04, -1.1521e-03, -4.3947e-03, -8.5493e-05, -2.7435e-05,\n        -1.4064e-04, -4.2322e-03,  3.0396e-03, -7.7105e-04, -1.2111e-03,\n        -9.3919e-04, -1.1787e-03, -1.9423e-03,  1.0555e-03,  1.9220e-03,\n        -8.3652e-04,  3.1769e-03,  1.5137e-03, -1.9903e-03, -1.4570e-03,\n         1.5441e-03,  3.2713e-03,  1.6705e-03,  1.7195e-03,  1.9656e-03,\n        -4.1300e-04,  1.9559e-03, -5.4321e-03, -1.5669e-03, -2.2334e-03,\n         2.4268e-03,  2.7842e-03,  4.8676e-03,  6.5762e-04, -8.7739e-04,\n        -6.4392e-03, -7.7498e-04, -1.7475e-03,  1.0667e-03, -1.3875e-03,\n        -5.3889e-04,  7.3358e-04, -1.1321e-03, -3.8308e-04, -9.5364e-04,\n         3.0327e-04,  1.7532e-03,  7.7406e-04, -6.5761e-04, -1.4295e-03,\n         2.8804e-03, -3.4327e-04,  1.2974e-03,  1.2066e-03,  6.5669e-04,\n         2.3948e-03,  1.8412e-03, -1.0351e-02,  3.7319e-03, -1.2403e-03,\n         4.8896e-04, -1.5015e-04, -2.0977e-03, -5.2152e-04, -2.4023e-04,\n         5.4263e-04, -1.2562e-02,  2.0852e-03,  2.6810e-03,  1.6668e-04,\n        -9.7694e-04, -1.5307e-03,  9.7498e-04, -5.8696e-05,  3.7297e-04,\n        -2.0433e-04, -5.3733e-03, -4.4945e-04,  1.9818e-03,  1.4138e-03,\n         8.7984e-04,  3.6068e-04,  1.0165e-03, -1.0187e-03,  5.3284e-04,\n         1.2509e-03,  1.3534e-03, -2.5382e-03, -4.6449e-03, -1.1655e-03,\n         2.8933e-04, -5.1909e-03,  1.3201e-03, -8.8282e-04,  4.1278e-04,\n        -1.8642e-03, -4.7300e-04,  4.0012e-04,  3.6222e-04,  2.4341e-03,\n         7.5427e-04, -1.9912e-03, -3.2787e-03,  1.2163e-03, -4.0500e-04,\n        -9.6603e-04, -1.7266e-03,  7.8661e-04,  4.0949e-03, -5.3389e-03,\n         3.5580e-03,  3.0716e-03,  3.2913e-03,  4.7870e-04,  9.9273e-04,\n        -3.3544e-03,  1.8581e-03, -9.1055e-04,  4.8808e-04, -4.0986e-04,\n        -3.3548e-03,  7.5316e-04, -3.1205e-03,  1.4747e-03,  8.9531e-04,\n         1.1409e-03, -2.4975e-03,  1.9354e-03,  1.7225e-04,  1.4483e-03,\n        -8.2022e-04,  1.0391e-03,  2.7731e-03,  8.6482e-04,  1.5326e-03,\n        -3.7473e-03, -3.1354e-04,  3.3995e-03, -1.8105e-04, -2.8787e-04,\n        -8.1246e-03, -3.2545e-03,  2.0741e-03, -1.9482e-03, -1.4571e-03,\n        -6.8159e-04,  7.9098e-04, -2.5212e-03, -1.1660e-03, -6.3629e-04,\n         2.1819e-03,  1.2585e-03,  3.1514e-04,  6.4672e-04,  4.3465e-03,\n         1.0540e-03, -3.3350e-03,  1.9752e-03, -4.5015e-04,  1.2544e-03,\n         3.8538e-03, -1.2945e-03,  2.2922e-03,  1.9641e-03, -1.4458e-03,\n         2.9125e-03, -5.2555e-03, -1.2550e-03, -4.5086e-04,  1.8257e-03,\n         1.6713e-03, -2.5974e-03, -1.5011e-03,  1.2974e-03, -5.5984e-04,\n         7.4047e-04, -4.3050e-03,  2.2720e-03, -1.1104e-04,  4.3742e-04,\n        -9.5901e-04,  7.1693e-04, -1.3201e-03,  4.0388e-04, -3.2207e-03,\n        -5.4032e-03,  9.7979e-04, -1.6721e-03, -1.9822e-03, -1.7101e-03,\n        -1.4230e-03,  1.6750e-03, -2.8188e-03, -1.2123e-03,  8.2985e-04,\n        -3.1523e-04, -9.4855e-04, -6.0235e-04, -2.6111e-03,  3.0811e-03,\n        -4.4801e-03, -9.1183e-04, -3.7965e-03, -7.7872e-04,  1.7078e-04,\n        -8.8201e-04,  1.0211e-03,  1.7733e-03, -2.4294e-03,  3.0546e-04,\n         1.1347e-03,  2.8915e-03,  6.0442e-04, -2.1089e-04, -1.5874e-04,\n        -5.0392e-04,  1.5891e-04,  3.3765e-03,  1.1523e-03,  2.3523e-03,\n         7.5852e-05,  2.2568e-03, -1.8872e-03, -5.0721e-03, -5.8292e-04,\n         3.1584e-03, -3.1155e-03,  1.5505e-03,  1.3268e-03,  4.9919e-04,\n         9.5487e-04, -2.1924e-03, -5.9790e-04, -1.1957e-03,  7.5898e-04,\n        -9.5407e-04,  4.0625e-04,  7.3509e-04, -7.1221e-04, -7.7495e-04,\n         5.2254e-04, -1.9950e-04, -8.1209e-03,  2.6934e-04, -1.7418e-03,\n        -6.0392e-04,  3.3616e-03, -1.0331e-03, -9.3736e-04,  2.6034e-03,\n         1.9261e-03,  2.2705e-03,  4.3738e-03, -3.5279e-03,  1.5247e-03,\n         1.2727e-03,  5.8993e-05, -4.4575e-03,  4.1534e-04, -3.7232e-04,\n        -3.6865e-06,  1.8960e-04, -7.5477e-04,  3.0180e-03, -2.6332e-04,\n         2.7126e-03,  3.2287e-03, -5.6548e-03, -8.4408e-04, -7.5098e-04,\n         3.8546e-03, -2.0561e-03,  4.3892e-03, -3.7621e-04, -6.6576e-04,\n         1.6865e-03,  2.4908e-03,  2.0277e-03,  1.8386e-03, -1.8361e-03,\n        -4.7032e-03,  2.9726e-03, -1.8173e-03, -2.0728e-03, -3.8943e-04,\n        -1.3252e-03,  1.1757e-03,  6.1412e-04, -1.5215e-03,  2.2864e-03,\n        -7.2901e-04, -1.1392e-04, -1.9575e-03,  2.4979e-03,  3.3975e-04,\n         6.5262e-04, -2.4489e-03,  1.3481e-03,  2.6548e-03,  1.7274e-05,\n         1.0811e-03, -2.7577e-04,  5.4188e-03,  1.6739e-04,  2.0993e-03,\n        -4.7362e-03, -6.7126e-04, -5.8165e-04,  4.8723e-04, -2.2167e-03,\n         3.4872e-04, -1.3695e-03,  4.6983e-04,  5.5412e-04,  5.1113e-04,\n        -5.1817e-03, -1.3361e-04,  2.3678e-03, -9.3557e-04,  8.6378e-04,\n         8.7715e-04, -3.7631e-03,  6.0648e-04,  1.5008e-03,  2.2842e-03,\n         7.0679e-04,  1.1943e-03, -1.2411e-03, -8.6336e-05,  2.3234e-03,\n        -4.5270e-06, -3.4252e-03, -3.8450e-04, -2.6986e-03, -1.1898e-03,\n         1.0207e-03,  5.0271e-04,  1.5110e-03, -6.0521e-03,  3.2933e-04,\n        -3.5102e-03,  3.3018e-05, -1.4676e-04, -1.1211e-03, -1.0126e-03,\n         1.2445e-03, -1.2805e-03, -2.9231e-03,  1.6986e-04, -6.8848e-03,\n        -2.9526e-03, -1.4090e-03, -3.0923e-03,  7.3828e-04, -6.2815e-04,\n         1.6140e-03,  5.2493e-04, -2.2259e-03,  2.9834e-04,  1.0423e-03,\n        -1.2359e-03,  5.9277e-03, -4.0567e-04, -2.2839e-03,  2.2975e-03,\n         8.5013e-04,  1.2484e-03,  9.4312e-04, -3.0186e-03, -1.1297e-03,\n        -6.7248e-03, -4.6395e-04,  2.6604e-03, -9.3869e-05,  2.6552e-03,\n        -5.0892e-04, -1.8761e-04,  2.3289e-03, -1.3559e-03, -3.5754e-04,\n        -1.7077e-03,  4.8932e-03,  1.9159e-03, -4.0395e-03,  1.0994e-03,\n        -7.8532e-04,  4.4951e-03], device='cuda:0')",
+        "exp_avg_sq": "tensor([3.6805e-05, 7.3937e-05, 5.8542e-05, 4.8088e-05, 4.3990e-05, 4.0886e-05,\n        2.6812e-05, 3.4669e-05, 5.0843e-05, 3.8308e-05, 4.7279e-05, 5.2565e-05,\n        6.5203e-05, 8.4050e-05, 5.7575e-05, 3.8645e-05, 6.9725e-05, 5.1133e-05,\n        6.2768e-05, 5.3528e-05, 5.3246e-05, 4.0388e-05, 5.1934e-05, 2.9991e-05,\n        1.0893e-04, 3.7622e-05, 5.0818e-05, 7.2045e-05, 6.4861e-05, 5.5137e-05,\n        4.9805e-05, 6.3562e-05, 3.0306e-05, 3.3061e-05, 7.6479e-05, 2.9636e-05,\n        6.9661e-05, 7.5459e-05, 2.9736e-05, 5.0717e-05, 3.8693e-05, 4.8324e-05,\n        4.2262e-05, 4.8798e-05, 3.5006e-05, 4.1160e-05, 4.8768e-05, 4.2778e-05,\n        5.9573e-05, 3.2704e-05, 5.4468e-05, 6.7364e-05, 4.3102e-05, 5.5270e-05,\n        4.2642e-05, 4.3430e-05, 3.8092e-05, 4.9398e-05, 5.3931e-05, 7.0820e-05,\n        6.1443e-05, 4.6501e-05, 5.6354e-05, 1.5672e-04, 4.8840e-05, 4.5152e-05,\n        6.3199e-05, 5.4005e-05, 3.1149e-05, 6.2683e-05, 5.1026e-05, 6.7359e-05,\n        5.0650e-05, 4.1484e-05, 3.3168e-05, 9.6477e-05, 4.6257e-05, 8.8566e-05,\n        5.0603e-05, 5.7138e-05, 6.5175e-05, 8.0373e-05, 8.1385e-05, 5.9455e-05,\n        6.2905e-05, 4.8052e-05, 6.2945e-05, 8.1142e-05, 6.4220e-11, 7.1123e-05,\n        4.5701e-05, 9.2766e-05, 4.2083e-05, 4.4945e-05, 8.5200e-05, 4.2624e-05,\n        4.2242e-05, 5.2274e-05, 4.8477e-05, 6.2582e-05, 4.0729e-05, 4.0337e-05,\n        3.6264e-05, 8.7707e-05, 5.6182e-05, 4.7374e-05, 4.4159e-05, 4.1929e-05,\n        3.7908e-05, 5.1894e-05, 4.5529e-05, 3.7234e-05, 5.9186e-05, 5.2459e-05,\n        4.2729e-05, 3.9106e-05, 5.7394e-05, 5.0099e-05, 3.9905e-05, 2.9276e-05,\n        5.7881e-05, 4.7650e-05, 5.6506e-05, 5.1832e-05, 3.9178e-05, 3.9626e-05,\n        3.1517e-05, 3.0676e-05, 4.3402e-05, 4.4518e-05, 3.8888e-05, 6.8498e-05,\n        5.4844e-05, 6.6723e-05, 8.1918e-05, 6.4189e-05, 4.3603e-05, 5.6950e-05,\n        3.8120e-05, 4.8239e-05, 5.7051e-05, 4.6541e-05, 5.7968e-05, 5.7197e-05,\n        5.3428e-05, 6.0866e-05, 7.3385e-05, 5.3849e-05, 4.7500e-05, 6.7541e-05,\n        3.5680e-05, 4.1777e-05, 4.6011e-05, 4.2314e-05, 4.3536e-05, 4.8088e-05,\n        5.4245e-05, 2.9563e-05, 6.1623e-05, 4.8354e-05, 4.7403e-05, 3.3478e-05,\n        4.7706e-05, 5.0560e-05, 2.5178e-05, 5.8178e-05, 4.8551e-05, 6.4073e-05,\n        4.9725e-05, 6.1510e-05, 4.1554e-05, 6.2202e-05, 5.8705e-05, 6.1497e-05,\n        8.6059e-05, 4.6032e-05, 5.1410e-05, 6.2621e-05, 4.8253e-05, 8.4380e-05,\n        7.0384e-05, 6.1315e-05, 5.5112e-05, 1.0852e-04, 7.6373e-05, 3.4335e-05,\n        6.9645e-05, 8.8747e-05, 4.7956e-05, 6.3612e-05, 4.5646e-05, 9.6595e-05,\n        6.8487e-05, 3.4970e-05, 5.8962e-05, 3.7908e-05, 4.8100e-05, 3.3616e-05,\n        4.5076e-05, 3.2201e-05, 6.0276e-05, 5.6248e-05, 5.7944e-05, 5.7331e-05,\n        4.0503e-05, 5.3161e-05, 3.5984e-05, 3.6681e-05, 3.2959e-05, 5.6431e-05,\n        3.8415e-05, 6.7970e-05, 7.1759e-05, 4.2055e-05, 7.3042e-05, 4.5809e-05,\n        4.5942e-05, 5.8677e-05, 6.2726e-05, 5.6845e-05, 6.5308e-05, 8.0984e-05,\n        6.9033e-05, 3.9029e-05, 5.2816e-05, 3.6832e-05, 7.3528e-05, 3.9509e-05,\n        4.5968e-05, 4.2868e-05, 2.6907e-05, 4.7027e-05, 5.4151e-05, 5.9676e-05,\n        6.3634e-05, 4.2313e-05, 8.9617e-05, 8.6283e-05, 3.4269e-05, 3.8427e-05,\n        5.9826e-05, 3.3533e-05, 3.3339e-05, 5.1069e-05, 3.5340e-05, 4.9429e-05,\n        3.5866e-05, 4.8919e-05, 6.0242e-05, 4.0861e-05, 5.5402e-05, 4.6662e-05,\n        6.1038e-05, 4.1311e-05, 4.6646e-05, 4.1644e-05, 3.2148e-05, 7.7027e-05,\n        3.5975e-05, 4.7480e-05, 4.9559e-05, 6.6777e-05, 6.3208e-05, 3.8730e-05,\n        4.1173e-05, 7.3003e-05, 4.1849e-05, 3.8839e-05, 3.9126e-05, 6.9553e-05,\n        4.0585e-05, 2.9138e-05, 6.2635e-05, 3.7909e-05, 4.0113e-05, 7.8943e-05,\n        5.4639e-05, 3.9306e-05, 6.6166e-05, 5.5038e-05, 5.9273e-05, 5.1681e-05,\n        9.0506e-05, 6.3974e-05, 4.2866e-05, 6.5564e-05, 4.0217e-05, 6.9711e-05,\n        6.2263e-05, 4.4330e-05, 3.6692e-05, 1.0448e-04, 4.6788e-05, 3.8016e-05,\n        4.7864e-05, 4.4140e-05, 9.7484e-05, 1.0209e-04, 2.2831e-05, 6.1538e-05,\n        3.8199e-05, 8.3140e-05, 3.7347e-05, 3.9100e-05, 5.4662e-05, 7.6378e-05,\n        6.5059e-05, 4.4603e-05, 4.1620e-05, 4.7557e-05, 5.5709e-05, 4.5426e-05,\n        3.7046e-05, 4.4605e-05, 2.6834e-05, 3.0995e-05, 6.3699e-05, 5.4689e-05,\n        4.8254e-05, 4.9489e-05, 5.9861e-05, 8.7019e-05, 8.2105e-05, 7.5413e-05,\n        6.8942e-05, 3.4976e-05, 4.0872e-05, 4.8758e-05, 4.2230e-05, 5.7723e-05,\n        5.3527e-05, 7.1423e-05, 5.4537e-05, 4.8284e-05, 3.0305e-05, 4.6845e-05,\n        4.7584e-05, 5.4110e-05, 4.3925e-05, 5.6366e-05, 5.5120e-05, 5.4311e-05,\n        4.3438e-05, 6.1051e-05, 4.8467e-05, 8.8103e-05, 6.8654e-05, 5.1174e-05,\n        3.7494e-05, 7.2669e-05, 4.0038e-05, 6.6344e-05, 5.6343e-05, 4.5152e-05,\n        4.7583e-05, 7.0478e-05, 5.1679e-05, 5.4520e-05, 6.6571e-05, 3.1084e-05,\n        6.3296e-05, 2.6308e-05, 3.3439e-05, 5.3752e-05, 5.2618e-05, 2.7523e-05,\n        3.7038e-05, 8.8722e-05, 7.1056e-05, 3.8261e-05, 8.3993e-05, 4.2115e-05,\n        4.8095e-05, 2.9971e-05, 5.9695e-05, 1.2726e-04, 3.6123e-05, 8.2811e-05,\n        4.8126e-05, 4.2755e-05, 5.5365e-05, 4.5476e-05, 1.2404e-04, 6.6355e-05,\n        6.3814e-05, 7.2946e-05, 3.1705e-05, 6.6186e-05, 5.8457e-05, 8.2941e-05,\n        6.4748e-05, 6.5316e-05, 3.9171e-05, 4.4058e-05, 5.1926e-05, 4.7984e-05,\n        3.8446e-05, 5.4465e-05, 5.1647e-05, 3.1666e-05, 7.6487e-05, 6.4791e-05,\n        5.2313e-05, 3.0687e-05, 3.3275e-05, 9.8068e-05, 4.6403e-05, 4.3078e-05,\n        4.8287e-05, 4.6822e-05, 5.4165e-05, 5.4477e-05, 4.7440e-05, 4.6478e-05,\n        5.7160e-05, 3.8572e-05, 5.0871e-05, 7.1905e-05, 6.5379e-05, 8.3975e-05,\n        4.8659e-05, 7.8890e-05, 5.9239e-05, 5.5878e-05, 5.8704e-05, 5.9433e-05,\n        4.2919e-05, 5.8995e-05, 4.5554e-05, 4.7780e-05, 3.8061e-05, 4.5313e-05,\n        3.7787e-05, 9.1572e-05, 4.3879e-05, 5.7953e-05, 8.3205e-05, 5.5552e-05,\n        5.6977e-05, 6.7190e-05, 1.2310e-04, 4.0099e-05, 4.8564e-05, 5.5910e-05,\n        4.8675e-05, 5.7628e-05, 5.4614e-05, 4.3573e-05, 5.9665e-05, 6.0839e-05,\n        8.8962e-05, 4.6997e-05, 4.5268e-05, 5.6266e-05, 5.0640e-05, 6.1417e-05,\n        3.4373e-05, 5.3757e-05, 5.2571e-05, 4.0398e-05, 6.0936e-05, 6.6020e-05,\n        5.2451e-05, 5.8451e-05, 7.0247e-05, 8.8714e-05, 6.4880e-05, 7.7108e-05,\n        6.8827e-05, 4.7016e-05, 9.8448e-05, 3.1812e-05, 2.9540e-05, 5.3065e-05,\n        5.1030e-05, 3.9378e-05, 4.2668e-05, 4.9246e-05, 3.8719e-05, 4.8421e-05,\n        4.6622e-05, 8.3083e-05, 6.2717e-05, 3.0664e-05, 3.8283e-05, 5.1384e-05,\n        6.8495e-05, 4.8804e-05, 5.5397e-05, 2.6952e-05, 6.3484e-05, 5.1081e-05,\n        4.0286e-05, 6.1025e-05, 4.4607e-05, 7.0251e-05, 3.9424e-05, 3.7349e-05,\n        3.9142e-05, 1.7464e-05, 6.8029e-05, 4.3056e-05, 3.9677e-05, 5.1420e-05,\n        3.7806e-05, 4.3370e-05, 5.8078e-05, 4.1708e-05, 8.6987e-05, 4.2758e-05,\n        5.5073e-05, 5.0651e-05], device='cuda:0')"
       },
       "4": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([[ 5.3741e-07, -1.8306e-05, -2.3914e-05,  ...,  2.7571e-05,\n         -7.6579e-06,  1.9829e-05],\n        [-2.7707e-05,  6.7255e-05,  6.8854e-07,  ..., -1.4027e-05,\n         -2.2276e-06,  2.5058e-05],\n        [ 2.2351e-05,  5.1829e-05, -2.9344e-06,  ..., -9.8238e-06,\n          9.9804e-06,  1.3086e-05],\n        ...,\n        [-2.6080e-05, -9.2305e-05,  1.9911e-05,  ...,  8.7396e-06,\n          6.5457e-06,  5.1776e-06],\n        [ 4.0660e-05, -1.8758e-05, -2.8677e-06,  ...,  5.6824e-06,\n         -1.4409e-05, -3.5914e-05],\n        [ 5.3487e-06, -3.5050e-05, -8.9959e-06,  ...,  8.2199e-06,\n          1.6605e-05, -2.2169e-05]], device='cuda:0')",
-        "exp_avg_sq": "tensor([[2.1125e-09, 4.1811e-09, 8.0170e-09,  ..., 3.0242e-09, 2.9669e-09,\n         6.9280e-09],\n        [4.7890e-09, 1.2046e-08, 6.0453e-09,  ..., 5.1771e-09, 3.7089e-09,\n         6.4596e-09],\n        [7.2451e-09, 1.1432e-08, 4.6009e-09,  ..., 7.5624e-09, 2.7876e-09,\n         1.2103e-08],\n        ...,\n        [4.4585e-09, 8.8713e-09, 6.1609e-09,  ..., 6.5394e-09, 3.1549e-09,\n         9.2782e-09],\n        [5.0723e-09, 7.6576e-09, 5.2430e-09,  ..., 6.1885e-09, 6.7636e-09,\n         1.1599e-08],\n        [5.6321e-09, 6.0320e-09, 8.2659e-09,  ..., 7.4554e-09, 4.0966e-09,\n         1.0538e-08]], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([[-5.9020e-06,  1.8948e-06,  1.6365e-05,  ...,  3.1623e-05,\n         -3.9503e-06,  1.3279e-06],\n        [ 8.9468e-06,  3.4703e-05,  9.5180e-06,  ...,  5.4239e-05,\n          1.9834e-05,  3.0988e-05],\n        [-1.1299e-05, -1.0186e-05, -1.4260e-05,  ..., -1.5816e-05,\n         -6.0828e-06,  1.3052e-05],\n        ...,\n        [ 9.0819e-06, -2.3612e-05,  3.1079e-05,  ...,  1.8375e-05,\n          2.3136e-07,  3.0862e-06],\n        [ 2.7104e-06, -2.3923e-06,  1.6058e-05,  ..., -1.9826e-05,\n         -1.3109e-05,  1.1551e-05],\n        [ 2.2167e-06, -1.8830e-05, -4.2202e-05,  ..., -1.0402e-05,\n         -2.1073e-05, -8.1033e-06]], device='cuda:0')",
+        "exp_avg_sq": "tensor([[2.0372e-09, 4.1398e-09, 7.6519e-09,  ..., 3.2138e-09, 3.1756e-09,\n         6.5789e-09],\n        [4.5471e-09, 1.1133e-08, 5.7402e-09,  ..., 4.9883e-09, 3.9634e-09,\n         6.0490e-09],\n        [6.5539e-09, 1.0672e-08, 4.5197e-09,  ..., 7.3073e-09, 2.8495e-09,\n         1.2085e-08],\n        ...,\n        [4.5090e-09, 8.3606e-09, 6.3867e-09,  ..., 6.4894e-09, 3.2104e-09,\n         9.0154e-09],\n        [5.1814e-09, 8.1224e-09, 5.4082e-09,  ..., 5.9186e-09, 7.4052e-09,\n         1.1339e-08],\n        [5.1179e-09, 6.1387e-09, 8.3817e-09,  ..., 7.3411e-09, 4.1967e-09,\n         1.0435e-08]], device='cuda:0')"
       },
       "5": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([[ 1.5232e-05, -1.2848e-05, -1.3202e-05,  ...,  1.4028e-05,\n         -8.4242e-06,  1.2926e-05],\n        [-1.6815e-06,  3.6981e-05, -2.6574e-06,  ..., -9.2885e-06,\n          4.1785e-06,  9.8522e-06],\n        [ 5.9082e-06,  5.5152e-05, -1.0712e-07,  ...,  5.6912e-06,\n          1.3168e-05,  1.2928e-06],\n        ...,\n        [-1.3415e-06,  6.0054e-05, -1.1864e-05,  ...,  3.7099e-06,\n         -6.2773e-06,  1.2221e-06],\n        [ 3.2136e-05,  7.3700e-05,  6.2024e-06,  ..., -1.1284e-05,\n          2.2259e-06,  6.0325e-06],\n        [ 7.0434e-06,  2.3958e-05, -5.4778e-06,  ...,  1.7597e-06,\n         -5.2584e-06, -4.5502e-06]], device='cuda:0')",
-        "exp_avg_sq": "tensor([[8.4779e-10, 1.3620e-09, 1.7582e-09,  ..., 1.2647e-09, 6.2579e-10,\n         2.3736e-09],\n        [1.8232e-09, 4.4844e-09, 1.6000e-09,  ..., 1.7763e-09, 1.9243e-09,\n         2.5399e-09],\n        [2.3894e-09, 5.4475e-09, 1.9017e-09,  ..., 3.0220e-09, 1.2430e-09,\n         4.0230e-09],\n        ...,\n        [2.5881e-09, 4.0688e-09, 2.2479e-09,  ..., 1.6522e-09, 1.6561e-09,\n         1.9196e-09],\n        [2.9269e-09, 7.8975e-09, 2.1694e-09,  ..., 2.2336e-09, 1.5162e-09,\n         3.1295e-09],\n        [1.3499e-09, 3.3986e-09, 1.5691e-09,  ..., 2.4193e-09, 1.1371e-09,\n         3.3536e-09]], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([[ 1.7989e-06,  4.2190e-06,  6.5448e-06,  ...,  3.0508e-06,\n         -1.8514e-06,  5.9576e-06],\n        [ 6.3176e-06,  1.6977e-05,  1.1505e-05,  ...,  3.7173e-05,\n          4.9809e-06,  1.5525e-05],\n        [-1.6708e-05,  9.5249e-06, -1.1249e-05,  ..., -1.5451e-05,\n         -7.9721e-06,  1.9833e-05],\n        ...,\n        [ 1.0364e-05,  1.1385e-05, -1.1192e-05,  ..., -8.4844e-06,\n         -5.6640e-06, -1.1401e-05],\n        [-1.7906e-06,  1.9834e-06,  1.0815e-06,  ..., -1.2380e-05,\n         -6.2534e-06, -9.2382e-06],\n        [ 9.6819e-06, -7.3784e-06,  4.8784e-06,  ..., -1.4418e-05,\n         -9.4953e-06, -7.2564e-06]], device='cuda:0')",
+        "exp_avg_sq": "tensor([[8.1560e-10, 1.3475e-09, 1.6525e-09,  ..., 1.2488e-09, 5.9595e-10,\n         2.3552e-09],\n        [1.7554e-09, 4.1794e-09, 1.5715e-09,  ..., 1.7226e-09, 1.9271e-09,\n         2.3219e-09],\n        [2.2232e-09, 4.9447e-09, 1.8756e-09,  ..., 2.8584e-09, 1.2778e-09,\n         3.8880e-09],\n        ...,\n        [2.5199e-09, 3.9330e-09, 2.5599e-09,  ..., 1.6822e-09, 1.6663e-09,\n         1.8369e-09],\n        [2.6684e-09, 7.1982e-09, 2.0619e-09,  ..., 2.2440e-09, 1.4449e-09,\n         2.9132e-09],\n        [1.2986e-09, 3.4277e-09, 1.6512e-09,  ..., 2.6567e-09, 1.1022e-09,\n         3.1725e-09]], device='cuda:0')"
       },
       "6": {
-        "step": "tensor(11268.)",
-        "exp_avg": "tensor([ 0.0010, -0.0010], device='cuda:0')",
-        "exp_avg_sq": "tensor([5.6570e-06, 5.6570e-06], device='cuda:0')"
+        "step": "tensor(12520.)",
+        "exp_avg": "tensor([ 0.0002, -0.0002], device='cuda:0')",
+        "exp_avg_sq": "tensor([5.9675e-06, 5.9675e-06], device='cuda:0')"
       }
     },
     "param_groups": [
       {
-        "lr": 0.00024569294678237997,
+        "lr": 0.01,
         "name": "shared",
         "betas": [
           0.9,
@@ -64,7 +64,7 @@
         ]
       },
       {
-        "lr": 0.00024569294678237997,
+        "lr": 0.01,
         "name": "scale_256",
         "betas": [
           0.9,
@@ -85,7 +85,7 @@
         ]
       },
       {
-        "lr": 0.00024569294678237997,
+        "lr": 0.01,
         "name": "scale_512",
         "betas": [
           0.9,
@@ -106,7 +106,7 @@
         ]
       },
       {
-        "lr": 0.00012333423752026375,
+        "lr": 0.005,
         "name": "fusion",
         "betas": [
           0.9,
@@ -130,33 +130,33 @@
   },
   "scheduler_state_dict": {
     "T_0": 10,
-    "T_i": 10,
+    "T_i": 20,
     "T_mult": 2,
     "eta_min": 1e-06,
-    "T_cur": 9,
+    "T_cur": 0,
     "base_lrs": [
       0.01,
       0.01,
       0.01,
       0.005
     ],
-    "last_epoch": 9,
+    "last_epoch": 10,
     "_step_count": 0,
     "_is_initial": false,
     "_get_lr_called_within_step": false,
     "_last_lr": [
-      0.00024569294678237997,
-      0.00024569294678237997,
-      0.00024569294678237997,
-      0.00012333423752026375
+      0.01,
+      0.01,
+      0.01,
+      0.005
     ]
   },
   "metrics": {
-    "best_val_acc": 75.24,
-    "best_epoch": 8,
+    "best_val_acc": 75.38,
+    "best_epoch": 9,
     "scale_accuracies": {
-      "256": 75.016,
-      "512": 75.192
+      "256": 75.102,
+      "512": 75.302
     }
   },
   "train_config": {