outputs

<s> tensor(1, device='cuda:0')
▁Donald tensor(18935, device='cuda:0')
▁was tensor(471, device='cuda:0')
▁very tensor(1407, device='cuda:0')
▁good tensor(1781, device='cuda:0')
▁at tensor(472, device='cuda:0')
▁playing tensor(8743, device='cuda:0')
▁the tensor(278, device='cuda:0')
▁viol tensor(5537, device='cuda:0')
in tensor(262, device='cuda:0')
▁but tensor(541, device='cuda:0')
▁Matthew tensor(22292, device='cuda:0')
▁was tensor(471, device='cuda:0')
▁not tensor(451, device='cuda:0')
. tensor(29889, device='cuda:0')
▁Matthew tensor(22292, device='cuda:0')
▁gave tensor(4846, device='cuda:0')
▁a tensor(263, device='cuda:0')
▁st tensor(380, device='cuda:0')
unning tensor(27389, device='cuda:0')
▁concert tensor(13135, device='cuda:0')
▁performance tensor(4180, device='cuda:0')
. tensor(29889, device='cuda:0')
<s>
▁Donald
▁was
▁very
▁good
▁at
▁playing
▁the
▁viol
in
▁but
▁Matthew
▁was
▁not
.
▁Matthew
▁gave
▁a
▁st
unning
▁concert
▁performance
.
(tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.8906e-01, 2.1289e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.6797e-01, 6.9824e-02, 2.6172e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.8828e-02, 1.5137e-02, 3.1982e-02,  ..., 1.8164e-01,
           0.0000e+00, 0.0000e+00],
          [4.7852e-02, 1.6479e-02, 2.9663e-02,  ..., 2.2656e-01,
           2.5757e-02, 0.0000e+00],
          [1.0254e-01, 1.8677e-02, 3.5400e-02,  ..., 7.7637e-02,
           4.1504e-02, 8.0078e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [5.7422e-01, 4.2578e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [3.7695e-01, 4.0039e-01, 2.2266e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.3213e-02, 3.8818e-02, 4.3701e-02,  ..., 4.8828e-02,
           0.0000e+00, 0.0000e+00],
          [4.1992e-02, 3.6133e-02, 4.1992e-02,  ..., 4.6143e-02,
           4.6143e-02, 0.0000e+00],
          [2.1973e-02, 3.0029e-02, 4.5166e-02,  ..., 4.1748e-02,
           5.6641e-02, 5.6152e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.2500e-01, 3.7695e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [4.4922e-01, 4.1406e-01, 1.3477e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [7.2266e-02, 2.4536e-02, 7.4219e-02,  ..., 3.7109e-02,
           0.0000e+00, 0.0000e+00],
          [7.3242e-02, 2.5391e-02, 6.1035e-02,  ..., 4.0771e-02,
           3.9307e-02, 0.0000e+00],
          [1.4355e-01, 2.8076e-02, 3.3691e-02,  ..., 3.3203e-02,
           3.3936e-02, 5.5176e-02]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9062e-01, 1.1035e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.6797e-01, 2.0215e-01, 1.2891e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [8.7402e-02, 1.1182e-01, 8.8882e-04,  ..., 1.0938e-01,
           0.0000e+00, 0.0000e+00],
          [4.8340e-02, 1.0449e-01, 7.8583e-04,  ..., 1.4746e-01,
           6.9336e-02, 0.0000e+00],
          [3.8574e-02, 1.2012e-01, 1.3916e-02,  ..., 9.6191e-02,
           4.6387e-02, 3.8086e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [2.0898e-01, 7.8906e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.9141e-01, 5.6641e-01, 2.4316e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.1035e-01, 7.8613e-02, 4.1260e-02,  ..., 1.4954e-02,
           0.0000e+00, 0.0000e+00],
          [1.0498e-01, 6.5918e-02, 3.4424e-02,  ..., 1.6724e-02,
           1.0757e-03, 0.0000e+00],
          [7.8613e-02, 5.9082e-02, 2.4780e-02,  ..., 2.2705e-02,
           2.8992e-03, 6.3965e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [5.3125e-01, 4.6875e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [2.0996e-01, 5.6641e-01, 2.2461e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1016e-02, 2.3193e-02, 7.1289e-02,  ..., 6.5918e-02,
           0.0000e+00, 0.0000e+00],
          [4.2480e-02, 2.4902e-02, 6.4453e-02,  ..., 7.5684e-02,
           5.8838e-02, 0.0000e+00],
          [1.6968e-02, 1.5869e-02, 6.1768e-02,  ..., 5.0781e-02,
           7.6660e-02, 2.2949e-02]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.8672, 0.1309, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.6484, 0.2695, 0.0811,  ..., 0.0000, 0.0000, 0.0000],
          ...,
          [0.1445, 0.0217, 0.0159,  ..., 0.0396, 0.0000, 0.0000],
          [0.1230, 0.0168, 0.0105,  ..., 0.0461, 0.0359, 0.0000],
          [0.1162, 0.0133, 0.0161,  ..., 0.0162, 0.0167, 0.2910]],

         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.4336, 0.5664, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.4277, 0.0564, 0.5156,  ..., 0.0000, 0.0000, 0.0000],
          ...,
          [0.1650, 0.0087, 0.0225,  ..., 0.1836, 0.0000, 0.0000],
          [0.1494, 0.0061, 0.0320,  ..., 0.0806, 0.1641, 0.0000],
          [0.1426, 0.0194, 0.0206,  ..., 0.0078, 0.0065, 0.2031]],

         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.7578, 0.2432, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.6758, 0.1748, 0.1484,  ..., 0.0000, 0.0000, 0.0000],
          ...,
          [0.0859, 0.0272, 0.0284,  ..., 0.0464, 0.0000, 0.0000],
          [0.0796, 0.0248, 0.0249,  ..., 0.0471, 0.0552, 0.0000],
          [0.1035, 0.0233, 0.0226,  ..., 0.0278, 0.0302, 0.1777]],

         ...,

         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.7617, 0.2393, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.4609, 0.1494, 0.3887,  ..., 0.0000, 0.0000, 0.0000],
          ...,
          [0.1328, 0.0540, 0.0287,  ..., 0.0253, 0.0000, 0.0000],
          [0.1357, 0.0554, 0.0304,  ..., 0.0234, 0.0275, 0.0000],
          [0.0776, 0.0486, 0.0432,  ..., 0.0312, 0.0325, 0.0820]],

         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.7109, 0.2910, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.7188, 0.1475, 0.1328,  ..., 0.0000, 0.0000, 0.0000],
          ...,
          [0.2676, 0.0425, 0.0476,  ..., 0.0164, 0.0000, 0.0000],
          [0.2773, 0.0403, 0.0459,  ..., 0.0137, 0.0157, 0.0000],
          [0.2236, 0.0422, 0.0420,  ..., 0.0223, 0.0267, 0.1035]],

         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.5508, 0.4512, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
          [0.3652, 0.2910, 0.3438,  ..., 0.0000, 0.0000, 0.0000],
          ...,
          [0.1318, 0.0175, 0.0306,  ..., 0.0679, 0.0000, 0.0000],
          [0.1279, 0.0177, 0.0304,  ..., 0.0669, 0.0742, 0.0000],
          [0.1016, 0.0278, 0.0369,  ..., 0.0405, 0.0437, 0.0913]]]],
       device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6484e-01, 3.6377e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 2.1851e-02, 9.2773e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.6484e-01, 3.9101e-04, 1.8358e-05,  ..., 1.2390e-02,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 9.7752e-05, 5.1856e-06,  ..., 9.6436e-03,
           1.0132e-02, 0.0000e+00],
          [7.0312e-02, 5.4932e-02, 6.3477e-02,  ..., 1.8311e-02,
           1.3550e-02, 6.9824e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 2.9297e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6484e-01, 1.9531e-02, 1.5747e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4531e-01, 5.1270e-03, 1.5640e-03,  ..., 7.4768e-03,
           0.0000e+00, 0.0000e+00],
          [4.5312e-01, 3.2349e-03, 9.9945e-04,  ..., 5.1575e-03,
           4.2725e-03, 0.0000e+00],
          [4.9072e-02, 4.9072e-02, 5.4199e-02,  ..., 3.6377e-02,
           3.6865e-02, 4.8584e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 2.0752e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0625e-01, 6.6406e-02, 2.8809e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5703e-01, 1.3161e-04, 1.2219e-05,  ..., 1.6724e-02,
           0.0000e+00, 0.0000e+00],
          [4.4336e-01, 2.7895e-05, 2.9355e-06,  ..., 4.0771e-02,
           1.6846e-02, 0.0000e+00],
          [3.4424e-02, 5.1025e-02, 6.1523e-02,  ..., 2.3804e-02,
           1.9531e-02, 3.3936e-02]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3750e-01, 6.3965e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0625e-01, 6.4453e-02, 2.9419e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.0234e-01, 1.8768e-03, 8.0872e-04,  ..., 2.5146e-02,
           0.0000e+00, 0.0000e+00],
          [4.2773e-01, 1.1673e-03, 5.0354e-04,  ..., 2.3804e-02,
           1.2939e-02, 0.0000e+00],
          [1.1279e-01, 5.0537e-02, 4.1016e-02,  ..., 2.7832e-02,
           2.0508e-02, 1.1377e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4531e-01, 5.5664e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.4062e-01, 3.4180e-01, 1.6724e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2773e-01, 2.0790e-04, 5.1022e-05,  ..., 1.9653e-02,
           0.0000e+00, 0.0000e+00],
          [3.7695e-01, 7.6771e-05, 1.3769e-05,  ..., 1.2598e-01,
           4.4678e-02, 0.0000e+00],
          [6.2256e-02, 5.8594e-02, 6.4941e-02,  ..., 2.3438e-02,
           2.0508e-02, 6.1523e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.9531e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 6.6528e-03, 6.7444e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.6875e-01, 5.1498e-04, 5.6458e-04,  ..., 6.3782e-03,
           0.0000e+00, 0.0000e+00],
          [4.7266e-01, 1.1635e-04, 2.0409e-04,  ..., 5.1880e-03,
           4.9438e-03, 0.0000e+00],
          [1.0156e-01, 5.1514e-02, 4.3457e-02,  ..., 2.4780e-02,
           2.0386e-02, 1.0156e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.8281e-01, 1.1670e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2578e-01, 1.0864e-02, 6.2500e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9219e-01, 7.6294e-04, 8.3923e-04,  ..., 4.0039e-02,
           0.0000e+00, 0.0000e+00],
          [5.0781e-01, 4.2343e-04, 8.9645e-04,  ..., 1.0132e-02,
           1.2390e-02, 0.0000e+00],
          [1.3281e-01, 3.4668e-02, 3.5400e-02,  ..., 2.3682e-02,
           2.7832e-02, 1.3184e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.2278e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.6469e-03, 2.0752e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.1562e-01, 8.2016e-04, 7.7057e-04,  ..., 7.0190e-04,
           0.0000e+00, 0.0000e+00],
          [5.1562e-01, 8.0109e-04, 7.5150e-04,  ..., 7.2861e-04,
           8.5068e-04, 0.0000e+00],
          [1.6895e-01, 2.6001e-02, 2.4658e-02,  ..., 2.9053e-02,
           2.6489e-02, 1.6113e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 4.2725e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 1.9775e-02, 2.6611e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.0000e-01, 2.5940e-04, 1.1110e-04,  ..., 1.8433e-02,
           0.0000e+00, 0.0000e+00],
          [4.9805e-01, 3.1281e-04, 9.8228e-05,  ..., 1.4832e-02,
           5.1270e-03, 0.0000e+00],
          [1.3574e-01, 4.6875e-02, 5.2734e-02,  ..., 2.1729e-02,
           1.7578e-02, 1.3086e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 9.6436e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 6.2866e-03, 3.2471e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9609e-01, 1.0529e-03, 4.3030e-03,  ..., 3.7842e-03,
           0.0000e+00, 0.0000e+00],
          [5.0781e-01, 4.0436e-04, 3.1433e-03,  ..., 2.1210e-03,
           1.8158e-03, 0.0000e+00],
          [2.5195e-01, 1.9653e-02, 2.0996e-02,  ..., 1.3733e-02,
           1.1963e-02, 2.3828e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3750e-01, 6.2256e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.2812e-01, 7.5684e-02, 9.6680e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2969e-01, 4.6082e-03, 9.4604e-03,  ..., 9.4604e-03,
           0.0000e+00, 0.0000e+00],
          [4.5117e-01, 3.4485e-03, 6.1340e-03,  ..., 6.7444e-03,
           3.5553e-03, 0.0000e+00],
          [2.5586e-01, 2.4536e-02, 2.0020e-02,  ..., 1.6113e-02,
           1.4160e-02, 2.4121e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.6060e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 7.3612e-06, 2.0508e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.2344e-01, 4.6968e-05, 3.1233e-05,  ..., 2.2827e-02,
           0.0000e+00, 0.0000e+00],
          [5.3516e-01, 3.9041e-06, 9.3579e-06,  ..., 4.6015e-05,
           6.8359e-03, 0.0000e+00],
          [1.0059e-01, 3.9307e-02, 4.4189e-02,  ..., 2.9053e-02,
           3.4180e-02, 9.7168e-02]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 6.1035e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 4.2236e-02, 5.7373e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5117e-01, 1.6499e-04, 1.1349e-04,  ..., 1.4893e-02,
           0.0000e+00, 0.0000e+00],
          [3.8281e-01, 5.0306e-05, 2.9564e-05,  ..., 1.9629e-01,
           2.9297e-02, 0.0000e+00],
          [1.4648e-01, 4.1504e-02, 3.9551e-02,  ..., 2.4902e-02,
           2.8076e-02, 1.4160e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.0071e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 9.8877e-03, 1.7944e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1406e-01, 2.2430e-03, 5.7068e-03,  ..., 4.1260e-02,
           0.0000e+00, 0.0000e+00],
          [4.4922e-01, 1.5335e-03, 3.4180e-03,  ..., 1.4465e-02,
           1.6479e-02, 0.0000e+00],
          [2.4805e-01, 1.6357e-02, 1.1536e-02,  ..., 3.5889e-02,
           1.9409e-02, 2.2754e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.0190e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.2817e-02, 5.7068e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4922e-01, 7.3853e-03, 1.0742e-02,  ..., 3.4943e-03,
           0.0000e+00, 0.0000e+00],
          [4.5898e-01, 6.4087e-03, 7.7209e-03,  ..., 4.6082e-03,
           3.0212e-03, 0.0000e+00],
          [2.7539e-01, 3.6621e-02, 9.7656e-03,  ..., 2.0996e-02,
           1.3367e-02, 2.4414e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.6479e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 7.5684e-03, 1.0803e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.6484e-01, 4.8828e-04, 7.3242e-04,  ..., 3.0151e-02,
           0.0000e+00, 0.0000e+00],
          [4.6680e-01, 3.2616e-04, 6.1035e-04,  ..., 2.6611e-02,
           2.1851e-02, 0.0000e+00],
          [3.0469e-01, 6.6833e-03, 6.3782e-03,  ..., 1.0193e-02,
           9.8267e-03, 2.7930e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.6733e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [5.1562e-01, 4.6875e-01, 1.5442e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1992e-01, 7.4463e-03, 5.3711e-03,  ..., 1.5869e-02,
           0.0000e+00, 0.0000e+00],
          [4.1602e-01, 8.1177e-03, 2.5391e-02,  ..., 3.5889e-02,
           4.4861e-03, 0.0000e+00],
          [2.5195e-01, 2.2827e-02, 2.0508e-02,  ..., 1.9775e-02,
           1.3611e-02, 2.3633e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 6.8359e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2969e-01, 5.8594e-02, 1.0315e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5703e-01, 2.5635e-03, 3.3569e-03,  ..., 1.3245e-02,
           0.0000e+00, 0.0000e+00],
          [3.8281e-01, 7.4387e-04, 7.9346e-04,  ..., 1.7969e-01,
           8.5449e-03, 0.0000e+00],
          [2.1680e-01, 3.4668e-02, 3.0029e-02,  ..., 2.1973e-02,
           1.1963e-02, 2.0508e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.6328e-01, 1.3672e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9453e-01, 8.0566e-02, 2.6489e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2969e-01, 4.6921e-04, 4.6921e-04,  ..., 2.1240e-02,
           0.0000e+00, 0.0000e+00],
          [4.1406e-01, 3.1662e-04, 4.0627e-04,  ..., 2.7100e-02,
           1.5747e-02, 0.0000e+00],
          [2.4512e-01, 2.5146e-02, 2.0996e-02,  ..., 1.5625e-02,
           1.1963e-02, 2.3242e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.8677e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0625e-01, 8.6914e-02, 4.7913e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4141e-01, 7.8735e-03, 1.4771e-02,  ..., 2.0409e-04,
           0.0000e+00, 0.0000e+00],
          [4.5898e-01, 4.6997e-03, 8.2397e-03,  ..., 1.5564e-03,
           9.4223e-04, 0.0000e+00],
          [2.0801e-01, 4.2236e-02, 3.3936e-02,  ..., 2.5269e-02,
           1.6846e-02, 2.0508e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.1641e-01, 1.8359e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1016e-01, 3.3936e-02, 5.4688e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.8633e-01, 1.7776e-03, 3.0670e-03,  ..., 1.1780e-02,
           0.0000e+00, 0.0000e+00],
          [4.9609e-01, 1.9379e-03, 2.1667e-03,  ..., 5.9814e-03,
           4.0894e-03, 0.0000e+00],
          [2.0898e-01, 3.5400e-02, 2.1240e-02,  ..., 2.4292e-02,
           1.8188e-02, 2.0410e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.3550e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9453e-01, 5.8105e-02, 4.6387e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.1172e-01, 1.5488e-03, 1.2054e-03,  ..., 3.8910e-03,
           0.0000e+00, 0.0000e+00],
          [5.0781e-01, 3.6049e-04, 9.4986e-04,  ..., 3.4790e-03,
           4.3335e-03, 0.0000e+00],
          [2.0410e-01, 3.7598e-02, 3.6133e-02,  ..., 2.3438e-02,
           2.2217e-02, 1.9336e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.4832e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.9922e-01, 2.9102e-01, 7.8735e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.9258e-01, 1.2939e-02, 2.9755e-03,  ..., 1.7456e-02,
           0.0000e+00, 0.0000e+00],
          [3.7891e-01, 2.3193e-03, 8.9111e-03,  ..., 6.7383e-02,
           1.3367e-02, 0.0000e+00],
          [2.1094e-01, 3.2471e-02, 2.4902e-02,  ..., 2.6123e-02,
           1.7822e-02, 1.9824e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1406e-01, 8.6426e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.3828e-01, 1.9043e-01, 7.1289e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.5195e-01, 2.6978e-02, 1.6235e-02,  ..., 4.8340e-02,
           0.0000e+00, 0.0000e+00],
          [2.7148e-01, 2.6123e-02, 1.8066e-02,  ..., 4.7852e-02,
           8.7891e-02, 0.0000e+00],
          [2.8320e-01, 1.6724e-02, 1.8677e-02,  ..., 1.7578e-02,
           1.5015e-02, 2.5586e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.4766e-01, 1.5137e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 1.5869e-02, 3.2715e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.9453e-01, 2.8076e-03, 3.3417e-03,  ..., 8.7891e-02,
           0.0000e+00, 0.0000e+00],
          [3.9844e-01, 1.7624e-03, 2.3346e-03,  ..., 8.4473e-02,
           3.2227e-02, 0.0000e+00],
          [2.4023e-01, 1.3770e-01, 1.5076e-02,  ..., 1.0681e-02,
           9.3994e-03, 2.4121e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3359e-01, 6.5918e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.0078e-01, 1.4648e-01, 5.1514e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.5527e-01, 1.7334e-02, 1.3123e-02,  ..., 4.8584e-02,
           0.0000e+00, 0.0000e+00],
          [2.8906e-01, 7.3547e-03, 1.4282e-02,  ..., 3.9307e-02,
           2.4658e-02, 0.0000e+00],
          [2.5977e-01, 1.0132e-02, 1.5015e-02,  ..., 3.0029e-02,
           1.7822e-02, 2.4316e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.4404e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4922e-01, 3.8574e-02, 1.2695e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2969e-01, 9.3994e-03, 1.5747e-02,  ..., 5.3024e-04,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 4.7302e-03, 8.5449e-03,  ..., 1.2894e-03,
           9.7656e-04, 0.0000e+00],
          [2.2461e-01, 3.5889e-02, 2.1606e-02,  ..., 1.9653e-02,
           1.2146e-02, 2.2461e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.7090e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6484e-01, 1.9775e-02, 1.3550e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4727e-01, 1.3428e-03, 4.4250e-03,  ..., 9.3384e-03,
           0.0000e+00, 0.0000e+00],
          [4.5508e-01, 1.0910e-03, 1.8539e-03,  ..., 1.2329e-02,
           1.3489e-02, 0.0000e+00],
          [2.4902e-01, 1.3184e-02, 1.7700e-02,  ..., 1.8066e-02,
           1.6113e-02, 2.3340e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.8311e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 9.7656e-03, 1.7548e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9805e-01, 6.1646e-03, 4.0588e-03,  ..., 1.3351e-03,
           0.0000e+00, 0.0000e+00],
          [5.0781e-01, 1.2283e-03, 2.8229e-04,  ..., 1.0834e-03,
           1.9653e-02, 0.0000e+00],
          [1.7383e-01, 2.3804e-02, 3.1250e-02,  ..., 3.0396e-02,
           2.6855e-02, 1.6309e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 8.0490e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 1.2085e-02, 3.2959e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.9062e-01, 9.7656e-04, 4.1199e-03,  ..., 3.4912e-02,
           0.0000e+00, 0.0000e+00],
          [4.2969e-01, 2.6321e-04, 7.1335e-04,  ..., 1.6968e-02,
           2.2827e-02, 0.0000e+00],
          [2.2461e-01, 1.0132e-02, 1.9287e-02,  ..., 2.9907e-02,
           2.8931e-02, 2.0605e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.0986e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 2.9907e-02, 1.2146e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5117e-01, 6.7444e-03, 2.4414e-03,  ..., 1.0254e-02,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 1.8463e-03, 1.3504e-03,  ..., 7.4158e-03,
           1.6251e-03, 0.0000e+00],
          [2.9492e-01, 1.0620e-02, 9.5215e-03,  ..., 9.2163e-03,
           7.9956e-03, 2.7344e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 2.9785e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.0703e-01, 4.8828e-02, 2.4414e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5117e-01, 1.5793e-03, 1.2283e-03,  ..., 6.7383e-02,
           0.0000e+00, 0.0000e+00],
          [4.5703e-01, 5.8746e-04, 7.7820e-04,  ..., 2.7466e-02,
           6.4453e-02, 0.0000e+00],
          [3.0078e-01, 1.2268e-02, 9.2773e-03,  ..., 5.3711e-03,
           4.3945e-03, 2.8711e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.2327e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9062e-01, 5.9814e-02, 4.8828e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5742e-01, 9.5215e-03, 2.9907e-03,  ..., 3.3447e-02,
           0.0000e+00, 0.0000e+00],
          [4.0820e-01, 2.7008e-03, 2.5787e-03,  ..., 3.1738e-02,
           1.0376e-02, 0.0000e+00],
          [3.3984e-01, 3.1891e-03, 3.0823e-03,  ..., 2.7313e-03,
           1.6022e-03, 3.0859e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.4038e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 8.9722e-03, 1.8921e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.6719e-01, 2.0905e-03, 1.2817e-03,  ..., 1.6357e-02,
           0.0000e+00, 0.0000e+00],
          [4.3359e-01, 8.9264e-04, 3.8338e-04,  ..., 4.7607e-03,
           5.6458e-03, 0.0000e+00],
          [2.8125e-01, 7.7515e-03, 1.5625e-02,  ..., 1.0132e-02,
           1.3000e-02, 2.6172e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.6245e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.6479e-02, 3.1433e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.0781e-01, 2.7657e-04, 1.5736e-04,  ..., 4.3945e-03,
           0.0000e+00, 0.0000e+00],
          [5.0781e-01, 5.1498e-05, 2.5153e-05,  ..., 3.5095e-03,
           4.3640e-03, 0.0000e+00],
          [3.1250e-01, 8.6060e-03, 6.7139e-03,  ..., 6.7139e-03,
           6.4087e-03, 2.8906e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.3008e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.9297e-01, 2.0508e-01, 2.1973e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.2344e-01, 3.0708e-04, 2.1839e-04,  ..., 1.1353e-02,
           0.0000e+00, 0.0000e+00],
          [5.2734e-01, 4.2677e-05, 6.6280e-05,  ..., 7.0801e-03,
           3.0060e-03, 0.0000e+00],
          [1.5918e-01, 2.2705e-02, 3.3203e-02,  ..., 5.1270e-02,
           4.7119e-02, 1.4844e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.3428e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 8.5449e-04, 1.8921e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.8633e-01, 3.4523e-04, 8.5449e-04,  ..., 3.7079e-03,
           0.0000e+00, 0.0000e+00],
          [4.9023e-01, 9.4891e-05, 2.9182e-04,  ..., 3.5095e-03,
           3.0518e-03, 0.0000e+00],
          [3.3594e-01, 3.8757e-03, 3.3722e-03,  ..., 4.8218e-03,
           4.0588e-03, 2.8906e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.2939e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2188e-01, 7.3242e-02, 3.4332e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5166e-02, 1.6937e-03, 4.8218e-03,  ..., 4.6082e-03,
           0.0000e+00, 0.0000e+00],
          [1.2891e-01, 1.4343e-03, 3.5400e-03,  ..., 1.2024e-02,
           2.0142e-03, 0.0000e+00],
          [1.9434e-01, 2.5635e-02, 4.1016e-02,  ..., 2.1729e-02,
           1.5869e-02, 1.8359e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.2125e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 3.1128e-03, 3.1128e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.0410e-01, 3.7575e-04, 1.0223e-03,  ..., 1.4099e-02,
           0.0000e+00, 0.0000e+00],
          [2.2852e-01, 2.2697e-04, 3.7384e-04,  ..., 1.7090e-02,
           3.0029e-02, 0.0000e+00],
          [1.9434e-01, 1.6724e-02, 1.9653e-02,  ..., 2.8931e-02,
           2.9419e-02, 1.8262e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.5547e-01, 1.4551e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.6328e-01, 1.1719e-01, 1.7944e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.0078e-01, 7.7820e-04, 7.0953e-04,  ..., 1.3367e-02,
           0.0000e+00, 0.0000e+00],
          [2.3340e-01, 2.3746e-04, 5.0354e-04,  ..., 3.1738e-02,
           1.5015e-02, 0.0000e+00],
          [2.2266e-01, 2.3438e-02, 2.6001e-02,  ..., 1.2573e-02,
           1.3123e-02, 2.0605e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 7.0953e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 7.2021e-03, 3.2959e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4336e-01, 1.4687e-04, 2.8229e-04,  ..., 3.5889e-02,
           0.0000e+00, 0.0000e+00],
          [3.4375e-01, 3.9339e-05, 8.2970e-05,  ..., 2.2949e-01,
           2.5513e-02, 0.0000e+00],
          [2.2363e-01, 1.5015e-02, 2.9541e-02,  ..., 2.5879e-02,
           2.1973e-02, 2.0898e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.1047e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 3.4790e-03, 7.1411e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.2734e-01, 9.9945e-04, 1.6251e-03,  ..., 2.3956e-03,
           0.0000e+00, 0.0000e+00],
          [5.2734e-01, 1.8883e-04, 5.5790e-05,  ..., 2.7313e-03,
           1.9287e-02, 0.0000e+00],
          [3.0078e-01, 6.3171e-03, 1.2756e-02,  ..., 1.2390e-02,
           1.1108e-02, 2.7148e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.9722e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 3.0060e-03, 8.4229e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.6680e-01, 2.3041e-03, 1.6327e-03,  ..., 1.8433e-02,
           0.0000e+00, 0.0000e+00],
          [4.2383e-01, 1.1597e-03, 1.9989e-03,  ..., 2.0874e-02,
           6.8970e-03, 0.0000e+00],
          [3.3008e-01, 5.8594e-03, 4.5166e-03,  ..., 6.5613e-03,
           4.5166e-03, 2.9102e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 4.0039e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [2.2070e-01, 7.7734e-01, 3.5858e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.3555e-01, 4.8218e-03, 1.1826e-03,  ..., 2.4261e-03,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 5.7602e-04, 1.9264e-04,  ..., 3.8147e-03,
           9.7656e-04, 0.0000e+00],
          [2.8516e-01, 2.1484e-02, 8.1177e-03,  ..., 5.7678e-03,
           4.8523e-03, 2.6953e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.0294e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0625e-01, 8.6426e-02, 8.1787e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.3555e-01, 3.7384e-04, 2.7275e-04,  ..., 4.2480e-02,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 2.8229e-04, 2.2030e-04,  ..., 3.0151e-02,
           2.5757e-02, 0.0000e+00],
          [2.1582e-01, 1.5381e-02, 1.7212e-02,  ..., 3.3936e-02,
           2.3315e-02, 2.0703e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.5747e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 1.6479e-02, 9.3994e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.6328e-01, 2.7924e-03, 1.6937e-03,  ..., 4.1260e-02,
           0.0000e+00, 0.0000e+00],
          [3.6133e-01, 1.2146e-02, 4.6997e-03,  ..., 3.1982e-02,
           3.1494e-02, 0.0000e+00],
          [3.6328e-01, 5.3024e-04, 3.9864e-04,  ..., 9.9182e-04,
           1.5564e-03, 3.0859e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.4404e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9453e-01, 5.3223e-02, 5.3955e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5352e-01, 9.2773e-03, 4.3640e-03,  ..., 1.7334e-02,
           0.0000e+00, 0.0000e+00],
          [3.3398e-01, 5.0354e-03, 2.3499e-03,  ..., 7.6172e-02,
           9.5215e-03, 0.0000e+00],
          [2.7344e-01, 8.6060e-03, 1.4832e-02,  ..., 2.0752e-02,
           1.7334e-02, 2.5195e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 2.3560e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.6489e-02, 1.7319e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5742e-01, 2.6855e-03, 2.2278e-03,  ..., 2.5757e-02,
           0.0000e+00, 0.0000e+00],
          [3.9453e-01, 8.0872e-04, 5.2261e-04,  ..., 3.0640e-02,
           1.0681e-02, 0.0000e+00],
          [2.8906e-01, 9.3994e-03, 9.5825e-03,  ..., 8.9722e-03,
           8.3008e-03, 2.6367e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.6692e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 9.5215e-03, 4.2114e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.8672e-01, 3.7842e-03, 2.6398e-03,  ..., 1.8799e-02,
           0.0000e+00, 0.0000e+00],
          [3.6914e-01, 4.1504e-03, 3.0823e-03,  ..., 1.6357e-02,
           2.6367e-02, 0.0000e+00],
          [3.1641e-01, 3.2806e-03, 5.6458e-03,  ..., 8.1177e-03,
           8.6060e-03, 2.8516e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [4.5898e-01, 5.3906e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [3.0859e-01, 1.0254e-02, 6.7969e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.3535e-01, 2.7657e-04, 1.6785e-04,  ..., 5.1562e-01,
           0.0000e+00, 0.0000e+00],
          [1.3574e-01, 1.0395e-04, 1.1778e-04,  ..., 5.0049e-02,
           6.6797e-01, 0.0000e+00],
          [9.0820e-02, 4.4434e-02, 4.0283e-02,  ..., 3.4180e-02,
           3.6621e-02, 8.9355e-02]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2188e-01, 7.9102e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.3047e-01, 1.7773e-01, 9.2285e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.3203e-01, 7.0496e-03, 8.3618e-03,  ..., 2.0874e-02,
           0.0000e+00, 0.0000e+00],
          [3.1250e-01, 2.7466e-03, 3.3112e-03,  ..., 1.0498e-01,
           1.7090e-02, 0.0000e+00],
          [3.3008e-01, 8.7891e-03, 7.5378e-03,  ..., 4.4861e-03,
           5.1575e-03, 2.9102e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 3.2471e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.4766e-01, 1.3672e-01, 1.4465e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.8828e-01, 1.0834e-03, 5.7220e-04,  ..., 5.3711e-03,
           0.0000e+00, 0.0000e+00],
          [4.5898e-01, 4.9591e-04, 1.8311e-04,  ..., 5.2002e-02,
           1.2939e-02, 0.0000e+00],
          [1.5527e-01, 2.1362e-02, 3.0396e-02,  ..., 4.6875e-02,
           2.9419e-02, 1.5137e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.1780e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1797e-01, 6.8359e-02, 1.2634e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.6895e-01, 1.5503e-02, 4.0894e-03,  ..., 1.7700e-02,
           0.0000e+00, 0.0000e+00],
          [1.8750e-01, 1.9409e-02, 7.4463e-03,  ..., 2.4170e-02,
           1.3428e-02, 0.0000e+00],
          [3.1250e-01, 5.2795e-03, 6.0730e-03,  ..., 5.0354e-03,
           5.0354e-03, 2.9492e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 2.0630e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.7109e-01, 1.0352e-01, 2.5635e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.8320e-01, 2.3651e-03, 5.0049e-03,  ..., 5.2490e-03,
           0.0000e+00, 0.0000e+00],
          [3.4375e-01, 4.3640e-03, 5.6152e-03,  ..., 7.9346e-03,
           6.0730e-03, 0.0000e+00],
          [3.0664e-01, 4.7913e-03, 6.9885e-03,  ..., 9.7046e-03,
           7.7820e-03, 2.7539e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2969e-01, 6.8848e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.9297e-01, 1.6992e-01, 3.7598e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.5977e-01, 6.3171e-03, 6.5918e-03,  ..., 1.4465e-02,
           0.0000e+00, 0.0000e+00],
          [2.3047e-01, 2.0264e-02, 1.0376e-02,  ..., 1.0864e-02,
           7.1716e-03, 0.0000e+00],
          [3.3203e-01, 6.3782e-03, 6.6833e-03,  ..., 3.4790e-03,
           3.5858e-03, 2.9688e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 6.8054e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.1414e-02, 5.7678e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.9043e-01, 1.7395e-03, 8.5449e-03,  ..., 3.5248e-03,
           0.0000e+00, 0.0000e+00],
          [3.2617e-01, 5.3406e-03, 1.0620e-02,  ..., 6.1646e-03,
           5.2490e-03, 0.0000e+00],
          [3.3398e-01, 4.6387e-03, 8.3618e-03,  ..., 6.1340e-03,
           6.1340e-03, 2.8516e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 3.5667e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 5.7220e-05, 8.6975e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.0781e-01, 2.3937e-04, 2.9755e-04,  ..., 6.5231e-04,
           0.0000e+00, 0.0000e+00],
          [5.6250e-01, 6.1989e-05, 1.3411e-05,  ..., 4.6492e-06,
           3.4523e-04, 0.0000e+00],
          [3.4766e-01, 1.2512e-03, 4.1504e-03,  ..., 4.3640e-03,
           1.8768e-03, 2.9492e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.0681e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 1.4343e-02, 1.2451e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5898e-01, 8.8501e-04, 7.3242e-04,  ..., 9.2163e-03,
           0.0000e+00, 0.0000e+00],
          [4.5703e-01, 9.0790e-04, 7.7820e-04,  ..., 8.2397e-03,
           1.1230e-02, 0.0000e+00],
          [3.2617e-01, 4.6387e-03, 5.0964e-03,  ..., 7.8735e-03,
           5.6763e-03, 2.9102e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.3956e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 3.9062e-03, 6.3782e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2773e-01, 2.0752e-03, 1.6308e-04,  ..., 1.7776e-03,
           0.0000e+00, 0.0000e+00],
          [4.4531e-01, 2.2583e-03, 2.9755e-04,  ..., 5.9814e-03,
           1.2131e-03, 0.0000e+00],
          [3.3008e-01, 2.9144e-03, 6.8970e-03,  ..., 5.8899e-03,
           8.1787e-03, 2.9297e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.1106e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 1.6968e-02, 6.5918e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.9922e-01, 2.8381e-03, 1.4267e-03,  ..., 2.5146e-02,
           0.0000e+00, 0.0000e+00],
          [3.1055e-01, 2.7771e-03, 1.7090e-03,  ..., 9.8267e-03,
           2.8076e-02, 0.0000e+00],
          [2.8711e-01, 9.3994e-03, 1.2451e-02,  ..., 1.0986e-02,
           1.3062e-02, 2.5781e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.7944e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1797e-01, 4.0283e-02, 4.1016e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.7539e-01, 7.7438e-04, 9.6130e-04,  ..., 1.4062e-01,
           0.0000e+00, 0.0000e+00],
          [3.8672e-01, 5.2261e-04, 1.0071e-03,  ..., 6.4453e-02,
           2.2583e-02, 0.0000e+00],
          [3.0078e-01, 6.2561e-03, 1.0132e-02,  ..., 9.0942e-03,
           6.9885e-03, 2.7734e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.5400e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 3.5156e-02, 1.0376e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.5977e-01, 5.6763e-03, 5.3406e-03,  ..., 1.1292e-02,
           0.0000e+00, 0.0000e+00],
          [3.4766e-01, 3.7079e-03, 7.1411e-03,  ..., 2.1484e-02,
           3.4180e-02, 0.0000e+00],
          [2.5977e-01, 1.7944e-02, 2.0386e-02,  ..., 1.3977e-02,
           1.0559e-02, 2.4512e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.1253e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.8215e-04, 4.6539e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.4688e-01, 4.2200e-05, 6.0272e-04,  ..., 4.6921e-04,
           0.0000e+00, 0.0000e+00],
          [5.5859e-01, 2.8729e-05, 1.2875e-04,  ..., 4.2152e-04,
           2.7466e-03, 0.0000e+00],
          [3.7109e-01, 7.7438e-04, 5.8365e-04,  ..., 1.1215e-03,
           8.7738e-04, 3.0078e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.9368e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 2.0630e-02, 8.7891e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5938e-01, 3.4809e-05, 1.0109e-04,  ..., 2.5513e-02,
           0.0000e+00, 0.0000e+00],
          [3.6328e-01, 6.3419e-05, 1.3351e-04,  ..., 5.7617e-02,
           1.8799e-02, 0.0000e+00],
          [3.1836e-01, 2.6550e-03, 3.2959e-03,  ..., 8.0566e-03,
           7.7820e-03, 2.9297e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.5747e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.1851e-02, 3.5706e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.0820e-01, 4.3335e-03, 3.0823e-03,  ..., 2.5940e-03,
           0.0000e+00, 0.0000e+00],
          [4.2773e-01, 3.4790e-03, 8.5068e-04,  ..., 4.8218e-03,
           2.6245e-03, 0.0000e+00],
          [2.8320e-01, 1.2451e-02, 1.0193e-02,  ..., 1.3489e-02,
           1.1719e-02, 2.6562e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.7334e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2188e-01, 6.4941e-02, 1.2146e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [7.0801e-02, 2.8419e-04, 1.2589e-04,  ..., 1.8311e-02,
           0.0000e+00, 0.0000e+00],
          [9.8145e-02, 5.2261e-04, 2.8992e-04,  ..., 3.3447e-02,
           9.8267e-03, 0.0000e+00],
          [2.8516e-01, 9.4604e-03, 7.5073e-03,  ..., 1.4465e-02,
           1.4038e-02, 2.6562e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.8555e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 2.5391e-02, 5.3406e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.9102e-01, 1.0490e-04, 2.5177e-04,  ..., 1.3306e-02,
           0.0000e+00, 0.0000e+00],
          [3.8086e-01, 2.3365e-04, 3.5095e-04,  ..., 2.2583e-02,
           7.0496e-03, 0.0000e+00],
          [2.7930e-01, 1.8555e-02, 1.4893e-02,  ..., 9.8877e-03,
           9.2773e-03, 2.5977e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.5776e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.0681e-02, 3.7537e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5547e-01, 1.9073e-03, 9.5749e-04,  ..., 1.6968e-02,
           0.0000e+00, 0.0000e+00],
          [4.1211e-01, 2.5024e-03, 6.5231e-04,  ..., 1.1047e-02,
           4.9744e-03, 0.0000e+00],
          [2.9297e-01, 1.3000e-02, 1.2573e-02,  ..., 9.5215e-03,
           7.9956e-03, 2.7734e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.9531e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 3.4424e-02, 5.8899e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.9258e-01, 8.0490e-04, 1.6117e-04,  ..., 1.6602e-01,
           0.0000e+00, 0.0000e+00],
          [2.0508e-01, 1.5717e-03, 3.0327e-04,  ..., 3.0859e-01,
           9.9609e-02, 0.0000e+00],
          [3.0469e-01, 7.7209e-03, 6.3171e-03,  ..., 7.8735e-03,
           6.0120e-03, 2.8711e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.7832e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.2217e-02, 5.4321e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.2461e-01, 1.5378e-05, 2.4676e-05,  ..., 4.1016e-02,
           0.0000e+00, 0.0000e+00],
          [1.9922e-01, 1.8954e-05, 1.3053e-05,  ..., 3.6914e-01,
           2.8198e-02, 0.0000e+00],
          [2.8125e-01, 1.4099e-02, 7.5378e-03,  ..., 1.0803e-02,
           1.2085e-02, 2.5586e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.4414e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4141e-01, 4.6875e-02, 1.0986e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5352e-01, 2.1973e-03, 1.2512e-03,  ..., 6.1035e-02,
           0.0000e+00, 0.0000e+00],
          [3.7305e-01, 1.0757e-03, 4.2915e-04,  ..., 9.4727e-02,
           8.2031e-02, 0.0000e+00],
          [3.0664e-01, 5.4932e-03, 5.6763e-03,  ..., 8.4229e-03,
           7.6599e-03, 2.9297e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 2.5146e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 2.6245e-02, 1.7700e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.6523e-01, 4.3640e-03, 4.8523e-03,  ..., 1.3611e-02,
           0.0000e+00, 0.0000e+00],
          [3.6328e-01, 6.6833e-03, 6.0120e-03,  ..., 1.5869e-02,
           1.0620e-02, 0.0000e+00],
          [2.6367e-01, 2.8931e-02, 1.9043e-02,  ..., 1.0376e-02,
           6.8054e-03, 2.5195e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.9531e-01, 3.0664e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.6172e-01, 2.0996e-02, 2.1680e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.7539e-01, 3.4790e-03, 2.9297e-03,  ..., 2.7539e-01,
           0.0000e+00, 0.0000e+00],
          [2.7344e-01, 1.9989e-03, 1.7853e-03,  ..., 1.0449e-01,
           2.3047e-01, 0.0000e+00],
          [2.2266e-01, 2.9297e-02, 2.4536e-02,  ..., 1.6602e-02,
           1.3977e-02, 2.1191e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.7548e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 3.1586e-03, 4.7302e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5117e-01, 1.8387e-03, 2.2278e-03,  ..., 3.9062e-03,
           0.0000e+00, 0.0000e+00],
          [4.9805e-01, 1.5869e-03, 2.1057e-03,  ..., 1.1749e-03,
           5.5237e-03, 0.0000e+00],
          [3.4180e-01, 5.0659e-03, 3.4790e-03,  ..., 3.2196e-03,
           3.0670e-03, 3.0469e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.9302e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2188e-01, 6.1523e-02, 1.7578e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.7344e-01, 4.0436e-04, 1.8597e-04,  ..., 7.5195e-02,
           0.0000e+00, 0.0000e+00],
          [2.9492e-01, 3.9673e-04, 4.3678e-04,  ..., 8.3984e-02,
           7.6660e-02, 0.0000e+00],
          [3.2031e-01, 6.8970e-03, 6.9885e-03,  ..., 4.9438e-03,
           4.0588e-03, 2.9492e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.3945e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 3.3447e-02, 1.0681e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.3398e-01, 7.8964e-04, 1.0300e-03,  ..., 6.2988e-02,
           0.0000e+00, 0.0000e+00],
          [2.9102e-01, 4.1771e-04, 5.2643e-04,  ..., 2.3438e-01,
           8.3008e-02, 0.0000e+00],
          [2.4805e-01, 1.0742e-02, 1.4709e-02,  ..., 2.0264e-02,
           1.5137e-02, 2.4219e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 5.7373e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.4375e-01, 6.6895e-02, 8.7891e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.6895e-01, 1.1368e-03, 1.4648e-03,  ..., 1.0742e-01,
           0.0000e+00, 0.0000e+00],
          [1.7480e-01, 1.1520e-03, 1.6479e-03,  ..., 1.3477e-01,
           1.4355e-01, 0.0000e+00],
          [3.1250e-01, 5.6152e-03, 6.8970e-03,  ..., 8.1787e-03,
           7.1106e-03, 2.8711e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.6328e-01, 1.3867e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1797e-01, 2.5269e-02, 5.7617e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.6484e-01, 4.6692e-03, 2.5024e-03,  ..., 5.8594e-02,
           0.0000e+00, 0.0000e+00],
          [4.9414e-01, 4.6997e-03, 1.8921e-03,  ..., 1.1780e-02,
           1.7944e-02, 0.0000e+00],
          [3.3203e-01, 9.8267e-03, 3.2959e-03,  ..., 3.2959e-03,
           3.2959e-03, 2.9883e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.2634e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.7891e-01, 5.9082e-02, 6.1035e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2188e-01, 6.1951e-03, 3.4180e-03,  ..., 3.3691e-02,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 6.8359e-03, 2.4719e-03,  ..., 9.2773e-03,
           6.9580e-03, 0.0000e+00],
          [3.3789e-01, 2.0508e-02, 6.0730e-03,  ..., 3.2043e-03,
           1.6632e-03, 3.0078e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3750e-01, 6.3965e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 1.6846e-02, 2.0264e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2578e-01, 6.1951e-03, 6.1035e-03,  ..., 1.4099e-02,
           0.0000e+00, 0.0000e+00],
          [4.1992e-01, 6.8359e-03, 1.4801e-03,  ..., 3.7109e-02,
           1.2817e-02, 0.0000e+00],
          [3.3008e-01, 8.7891e-03, 2.5940e-03,  ..., 1.0834e-03,
           1.1520e-03, 3.1445e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.8931e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0234e-01, 6.1523e-02, 3.6133e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.5195e-01, 2.1973e-03, 1.5564e-03,  ..., 9.7168e-02,
           0.0000e+00, 0.0000e+00],
          [3.1836e-01, 1.5411e-03, 1.3580e-03,  ..., 1.1670e-01,
           1.3379e-01, 0.0000e+00],
          [2.4023e-01, 2.9663e-02, 2.1240e-02,  ..., 1.7090e-02,
           1.5015e-02, 2.2754e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6484e-01, 3.5889e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0234e-01, 6.1035e-02, 3.6377e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5352e-01, 8.8501e-03, 3.5095e-03,  ..., 5.5237e-03,
           0.0000e+00, 0.0000e+00],
          [3.6914e-01, 1.1169e-02, 5.6458e-03,  ..., 6.7139e-03,
           4.2725e-03, 0.0000e+00],
          [2.6172e-01, 1.9409e-02, 1.7700e-02,  ..., 2.1362e-02,
           1.5381e-02, 2.5195e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.4463e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 3.5156e-02, 7.4158e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.7461e-01, 6.7139e-04, 2.8801e-04,  ..., 1.2939e-02,
           0.0000e+00, 0.0000e+00],
          [4.3945e-01, 6.2180e-04, 5.1498e-04,  ..., 1.5747e-02,
           3.6133e-02, 0.0000e+00],
          [3.3594e-01, 2.8381e-03, 1.4954e-03,  ..., 5.5542e-03,
           5.0659e-03, 3.0859e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.6163e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4922e-01, 4.0527e-02, 1.0620e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1797e-01, 1.1444e-03, 3.6621e-04,  ..., 2.8564e-02,
           0.0000e+00, 0.0000e+00],
          [4.5898e-01, 9.8419e-04, 3.9864e-04,  ..., 7.9956e-03,
           1.0559e-02, 0.0000e+00],
          [3.3398e-01, 5.3406e-03, 3.6163e-03,  ..., 1.8768e-03,
           1.1063e-03, 3.1055e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 5.2185e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 1.7578e-02, 2.3682e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4727e-01, 7.2479e-04, 1.3580e-03,  ..., 2.9907e-02,
           0.0000e+00, 0.0000e+00],
          [4.6094e-01, 8.8882e-04, 1.7166e-03,  ..., 1.3916e-02,
           1.1475e-02, 0.0000e+00],
          [3.5742e-01, 8.0109e-04, 3.7956e-04,  ..., 2.1515e-03,
           1.7014e-03, 3.1445e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.7548e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 7.2098e-04, 2.2125e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.2461e-01, 2.3804e-03, 6.2180e-04,  ..., 8.1177e-03,
           0.0000e+00, 0.0000e+00],
          [2.3145e-01, 1.4191e-03, 3.1090e-04,  ..., 2.8687e-03,
           7.5378e-03, 0.0000e+00],
          [3.7109e-01, 1.9455e-03, 1.2131e-03,  ..., 9.7656e-04,
           1.1444e-03, 3.0469e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.3651e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.5547e-01, 4.2236e-02, 1.0254e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.0664e-01, 1.2283e-03, 9.9182e-04,  ..., 1.7578e-01,
           0.0000e+00, 0.0000e+00],
          [2.1289e-01, 6.0654e-04, 1.8387e-03,  ..., 3.4668e-02,
           4.2969e-01, 0.0000e+00],
          [3.3398e-01, 3.0975e-03, 4.1809e-03,  ..., 3.9062e-03,
           1.8234e-03, 3.0859e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 9.3994e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 3.5553e-03, 3.5553e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4531e-01, 2.6245e-03, 3.6469e-03,  ..., 8.9111e-03,
           0.0000e+00, 0.0000e+00],
          [4.6680e-01, 8.6212e-04, 2.0599e-03,  ..., 4.7913e-03,
           6.9885e-03, 0.0000e+00],
          [3.3594e-01, 3.9368e-03, 4.3335e-03,  ..., 3.5248e-03,
           2.8381e-03, 3.0469e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.0588e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 5.8899e-03, 2.5787e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.4219e-01, 7.7148e-02, 5.5420e-02,  ..., 1.8188e-02,
           0.0000e+00, 0.0000e+00],
          [2.0605e-01, 5.9082e-02, 2.9541e-02,  ..., 5.2490e-03,
           1.0300e-03, 0.0000e+00],
          [3.4375e-01, 2.6703e-03, 2.0447e-03,  ..., 1.3199e-03,
           1.0986e-03, 3.1641e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.9165e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1016e-01, 2.1484e-02, 6.9336e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5508e-01, 3.6049e-04, 8.3923e-04,  ..., 8.1543e-02,
           0.0000e+00, 0.0000e+00],
          [3.7891e-01, 1.5545e-04, 5.4550e-04,  ..., 7.2754e-02,
           1.3965e-01, 0.0000e+00],
          [3.1641e-01, 7.5989e-03, 7.1411e-03,  ..., 7.1411e-03,
           5.4016e-03, 2.9102e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.5400e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 8.2397e-03, 2.1729e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.3203e-01, 4.1748e-02, 3.3936e-02,  ..., 7.1716e-03,
           0.0000e+00, 0.0000e+00],
          [3.3008e-01, 3.0151e-02, 2.8442e-02,  ..., 4.6997e-03,
           3.0823e-03, 0.0000e+00],
          [3.0469e-01, 1.9165e-02, 1.7700e-02,  ..., 8.6060e-03,
           7.1411e-03, 2.5391e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 3.2471e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3359e-01, 2.6489e-02, 3.9795e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.7852e-01, 3.6774e-03, 7.2021e-03,  ..., 2.6733e-02,
           0.0000e+00, 0.0000e+00],
          [5.0391e-01, 2.4414e-03, 3.9673e-03,  ..., 2.1484e-02,
           8.7891e-03, 0.0000e+00],
          [3.5352e-01, 8.4229e-03, 6.8665e-03,  ..., 3.0060e-03,
           1.6556e-03, 2.8320e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.6235e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.0132e-02, 1.9073e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.6133e-01, 4.8523e-03, 1.9302e-03,  ..., 4.9438e-03,
           0.0000e+00, 0.0000e+00],
          [3.8281e-01, 3.8757e-03, 2.4567e-03,  ..., 5.3711e-03,
           4.6692e-03, 0.0000e+00],
          [3.6133e-01, 3.6621e-03, 3.9673e-03,  ..., 1.9684e-03,
           1.2283e-03, 2.8906e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [4.6680e-01, 5.3125e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1016e-01, 9.6436e-03, 8.0566e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.3125e-01, 3.9864e-04, 3.0136e-04,  ..., 2.7100e-02,
           0.0000e+00, 0.0000e+00],
          [5.5469e-01, 2.2602e-04, 1.1730e-04,  ..., 2.7924e-03,
           7.7209e-03, 0.0000e+00],
          [3.4180e-01, 1.3855e-02, 2.0905e-03,  ..., 2.3804e-03,
           2.0294e-03, 3.1055e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 9.3994e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0234e-01, 6.5918e-02, 3.0640e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.0820e-01, 1.8921e-03, 4.9973e-04,  ..., 7.5989e-03,
           0.0000e+00, 0.0000e+00],
          [3.5742e-01, 1.8158e-03, 5.9891e-04,  ..., 1.3977e-02,
           2.0264e-02, 0.0000e+00],
          [3.6719e-01, 1.7853e-03, 1.1139e-03,  ..., 1.9302e-03,
           1.1520e-03, 3.0078e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 6.1035e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 4.6692e-03, 1.7334e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.2422e-01, 4.0894e-03, 6.3477e-02,  ..., 2.7344e-02,
           0.0000e+00, 0.0000e+00],
          [2.8320e-01, 1.5335e-03, 4.4678e-02,  ..., 3.1494e-02,
           2.8198e-02, 0.0000e+00],
          [4.0625e-01, 1.4420e-03, 1.4038e-03,  ..., 1.8539e-03,
           1.1292e-03, 2.8125e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.6469e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 5.0049e-03, 9.5215e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.4160e-01, 4.3678e-04, 6.0654e-04,  ..., 1.8164e-01,
           0.0000e+00, 0.0000e+00],
          [1.4746e-01, 3.1281e-04, 9.3460e-04,  ..., 1.8359e-01,
           1.1572e-01, 0.0000e+00],
          [3.7109e-01, 1.4191e-03, 1.9455e-03,  ..., 2.7771e-03,
           3.0060e-03, 2.9297e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.1787e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 1.7334e-02, 1.3306e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.2227e-01, 7.3547e-03, 7.0190e-03,  ..., 3.8147e-03,
           0.0000e+00, 0.0000e+00],
          [3.4375e-01, 5.6458e-03, 6.9275e-03,  ..., 3.9978e-03,
           5.5542e-03, 0.0000e+00],
          [3.9648e-01, 3.0518e-03, 1.9379e-03,  ..., 7.1335e-04,
           7.1335e-04, 2.8711e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.3984e-01, 1.6113e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.4766e-01, 1.4453e-01, 8.7891e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.3145e-01, 3.9339e-05, 8.7619e-06,  ..., 5.1953e-01,
           0.0000e+00, 0.0000e+00],
          [3.6328e-01, 4.1008e-05, 1.3769e-05,  ..., 5.2979e-02,
           2.5781e-01, 0.0000e+00],
          [3.4375e-01, 1.6479e-02, 4.6997e-03,  ..., 5.6763e-03,
           5.6763e-03, 2.7930e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.2031e-01, 1.7969e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 9.0942e-03, 3.1006e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.4688e-01, 4.5967e-04, 8.5831e-04,  ..., 9.0332e-03,
           0.0000e+00, 0.0000e+00],
          [5.3906e-01, 3.8910e-04, 3.8910e-04,  ..., 5.3711e-03,
           5.7983e-03, 0.0000e+00],
          [3.7891e-01, 3.6926e-03, 3.4637e-03,  ..., 2.0752e-03,
           1.7700e-03, 2.9297e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.6602e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 2.2583e-02, 7.2327e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.8652e-01, 2.4609e-01, 2.7954e-02,  ..., 1.8845e-03,
           0.0000e+00, 0.0000e+00],
          [1.6895e-01, 1.5137e-01, 2.7100e-02,  ..., 1.9684e-03,
           2.7618e-03, 0.0000e+00],
          [3.5352e-01, 1.1108e-02, 8.4229e-03,  ..., 2.5635e-03,
           2.2583e-03, 2.8516e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 4.9591e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 7.3853e-03, 1.0071e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9414e-01, 3.3264e-03, 2.8839e-03,  ..., 3.1006e-02,
           0.0000e+00, 0.0000e+00],
          [4.1211e-01, 4.3030e-03, 4.1809e-03,  ..., 4.1992e-02,
           3.0273e-02, 0.0000e+00],
          [3.6914e-01, 2.6093e-03, 3.0975e-03,  ..., 3.4637e-03,
           2.3804e-03, 2.9297e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 5.7678e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 6.3477e-03, 6.2561e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5312e-01, 5.8746e-04, 3.6812e-04,  ..., 1.9775e-02,
           0.0000e+00, 0.0000e+00],
          [4.5703e-01, 1.0605e-03, 1.0910e-03,  ..., 9.5825e-03,
           8.7280e-03, 0.0000e+00],
          [3.7109e-01, 3.4790e-03, 1.9836e-03,  ..., 1.4954e-03,
           1.0300e-03, 2.9883e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.5513e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2578e-01, 4.6387e-02, 2.7710e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1016e-01, 1.5030e-03, 3.6812e-04,  ..., 2.5879e-02,
           0.0000e+00, 0.0000e+00],
          [3.8867e-01, 9.9945e-04, 3.2616e-04,  ..., 2.1484e-02,
           4.1992e-02, 0.0000e+00],
          [3.0273e-01, 1.1353e-02, 9.1553e-03,  ..., 7.3242e-03,
           5.6152e-03, 2.7539e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.8281e-01, 1.1670e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 4.9133e-03, 4.0894e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9609e-01, 5.7936e-05, 2.1338e-05,  ..., 1.3977e-02,
           0.0000e+00, 0.0000e+00],
          [4.9609e-01, 5.8889e-05, 1.6809e-05,  ..., 2.9541e-02,
           5.7983e-03, 0.0000e+00],
          [3.3008e-01, 5.2185e-03, 2.2888e-03,  ..., 2.1515e-03,
           1.4267e-03, 3.2617e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.1797e-01, 8.2031e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 5.5237e-03, 1.2634e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.9648e-01, 8.5449e-04, 1.5717e-03,  ..., 2.2949e-01,
           0.0000e+00, 0.0000e+00],
          [4.6484e-01, 6.9427e-04, 1.7700e-03,  ..., 3.4668e-02,
           1.0938e-01, 0.0000e+00],
          [3.1445e-01, 1.2390e-02, 4.2114e-03,  ..., 1.6235e-02,
           1.2390e-02, 2.7148e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.0190e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.4766e-01, 1.2573e-02, 1.4062e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.3516e-01, 1.9741e-04, 9.9945e-04,  ..., 3.1738e-03,
           0.0000e+00, 0.0000e+00],
          [5.3906e-01, 2.6703e-04, 1.3199e-03,  ..., 1.7166e-03,
           8.7738e-04, 0.0000e+00],
          [4.0430e-01, 5.9128e-04, 9.4604e-04,  ..., 3.2806e-04,
           2.1172e-04, 2.9102e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.7079e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.0625e-01, 5.7861e-02, 3.6621e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5117e-01, 6.8665e-04, 9.9182e-05,  ..., 4.0771e-02,
           0.0000e+00, 0.0000e+00],
          [4.6094e-01, 9.4986e-04, 1.1683e-04,  ..., 1.6235e-02,
           3.0884e-02, 0.0000e+00],
          [3.4180e-01, 3.7689e-03, 1.7242e-03,  ..., 3.8300e-03,
           2.1515e-03, 3.0664e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.4343e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 3.1433e-03, 9.3994e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1211e-01, 8.3008e-03, 1.3977e-02,  ..., 1.2741e-03,
           0.0000e+00, 0.0000e+00],
          [3.3789e-01, 1.4587e-02, 1.5198e-02,  ..., 4.8065e-04,
           1.1368e-03, 0.0000e+00],
          [4.1797e-01, 1.9989e-03, 5.0354e-03,  ..., 1.7624e-03,
           1.2131e-03, 2.7344e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 9.3842e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 8.0109e-04, 1.8387e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.1328e-01, 4.3335e-03, 1.4496e-03,  ..., 2.9297e-03,
           0.0000e+00, 0.0000e+00],
          [6.2500e-01, 5.4016e-03, 2.1820e-03,  ..., 1.3428e-03,
           4.1008e-04, 0.0000e+00],
          [4.2578e-01, 3.0365e-03, 1.8692e-03,  ..., 2.4719e-03,
           1.1673e-03, 2.6562e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.9297e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 1.0452e-03, 5.3711e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.5469e-01, 4.6539e-04, 8.2016e-04,  ..., 6.4392e-03,
           0.0000e+00, 0.0000e+00],
          [5.5859e-01, 3.8338e-04, 9.4986e-04,  ..., 6.2866e-03,
           1.9455e-03, 0.0000e+00],
          [4.3359e-01, 8.6975e-04, 7.6675e-04,  ..., 6.9809e-04,
           4.9591e-04, 2.7734e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.0834e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 3.4523e-04, 1.1368e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.1719e-01, 9.4223e-04, 1.7319e-03,  ..., 1.0834e-03,
           0.0000e+00, 0.0000e+00],
          [6.2500e-01, 8.5449e-04, 1.5182e-03,  ..., 3.3379e-04,
           1.1292e-03, 0.0000e+00],
          [4.2773e-01, 1.7853e-03, 2.4414e-03,  ..., 1.4343e-03,
           1.7853e-03, 2.7148e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.0823e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 5.3406e-04, 2.0447e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.0391e-01, 4.1504e-03, 1.8997e-03,  ..., 2.1240e-02,
           0.0000e+00, 0.0000e+00],
          [4.9414e-01, 4.2725e-03, 3.9673e-03,  ..., 4.8523e-03,
           3.3264e-03, 0.0000e+00],
          [4.3555e-01, 1.7319e-03, 1.9531e-03,  ..., 1.3885e-03,
           1.2665e-03, 2.6562e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.6703e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4531e-01, 2.7710e-02, 2.6855e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.2578e-01, 2.1210e-03, 1.8997e-03,  ..., 5.6152e-02,
           0.0000e+00, 0.0000e+00],
          [4.6875e-01, 2.6855e-03, 3.8452e-03,  ..., 3.5889e-02,
           2.5391e-02, 0.0000e+00],
          [4.0039e-01, 1.9836e-03, 3.3875e-03,  ..., 2.5177e-03,
           1.5488e-03, 2.7148e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.2146e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 8.1177e-03, 2.9297e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.5078e-01, 2.1729e-02, 1.5488e-03,  ..., 1.5015e-02,
           0.0000e+00, 0.0000e+00],
          [5.6250e-01, 1.8188e-02, 1.3199e-03,  ..., 9.5825e-03,
           8.4839e-03, 0.0000e+00],
          [3.8086e-01, 4.3030e-03, 2.7313e-03,  ..., 3.0060e-03,
           2.1973e-03, 2.9102e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.6060e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 4.4556e-03, 2.4261e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.0391e-01, 1.8433e-02, 5.9814e-03,  ..., 5.8365e-04,
           0.0000e+00, 0.0000e+00],
          [4.7656e-01, 1.5991e-02, 6.5002e-03,  ..., 1.9646e-04,
           1.0605e-03, 0.0000e+00],
          [4.0625e-01, 6.8054e-03, 4.1199e-03,  ..., 1.4267e-03,
           2.3499e-03, 2.6562e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.4332e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 1.6098e-03, 6.6833e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.3633e-01, 5.3406e-03, 8.6670e-03,  ..., 8.8379e-02,
           0.0000e+00, 0.0000e+00],
          [2.3242e-01, 3.9978e-03, 1.0681e-02,  ..., 8.6914e-02,
           6.0547e-02, 0.0000e+00],
          [4.0039e-01, 3.3722e-03, 1.5488e-03,  ..., 2.2430e-03,
           1.4496e-03, 2.7734e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 4.7852e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 5.3101e-03, 4.6349e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.1328e-01, 1.5259e-02, 5.4932e-04,  ..., 1.1826e-03,
           0.0000e+00, 0.0000e+00],
          [6.2891e-01, 1.3306e-02, 5.3024e-04,  ..., 5.6458e-04,
           1.3123e-03, 0.0000e+00],
          [4.1406e-01, 9.5825e-03, 2.0599e-03,  ..., 2.9144e-03,
           3.6316e-03, 2.6953e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.0190e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.1902e-03, 2.5749e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.4727e-01, 1.6016e-01, 4.1809e-03,  ..., 1.4648e-03,
           0.0000e+00, 0.0000e+00],
          [4.4727e-01, 1.1133e-01, 6.8970e-03,  ..., 1.4114e-04,
           5.5695e-04, 0.0000e+00],
          [4.4531e-01, 3.6163e-03, 2.3346e-03,  ..., 5.7220e-04,
           7.3242e-04, 2.5586e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.6846e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 1.4343e-02, 1.2268e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.4688e-01, 5.3406e-03, 8.8882e-04,  ..., 1.2741e-03,
           0.0000e+00, 0.0000e+00],
          [4.9609e-01, 1.0071e-02, 1.6022e-03,  ..., 4.6997e-03,
           8.7891e-03, 0.0000e+00],
          [3.6914e-01, 4.7913e-03, 2.1973e-03,  ..., 4.3030e-03,
           3.0975e-03, 2.8711e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 3.7354e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.8672e-01, 4.3030e-03, 1.0986e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.4688e-01, 1.7929e-04, 7.0953e-04,  ..., 8.2397e-03,
           0.0000e+00, 0.0000e+00],
          [5.5859e-01, 1.3733e-04, 6.9809e-04,  ..., 2.1210e-03,
           4.1504e-03, 0.0000e+00],
          [3.9648e-01, 2.5330e-03, 1.5411e-03,  ..., 6.5994e-04,
           4.6921e-04, 2.9297e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.2114e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.1641e-01, 5.5847e-03, 1.7773e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.5469e-01, 1.8597e-04, 5.0354e-04,  ..., 1.9989e-03,
           0.0000e+00, 0.0000e+00],
          [5.4688e-01, 2.2888e-04, 9.0408e-04,  ..., 1.9150e-03,
           3.8147e-03, 0.0000e+00],
          [3.9258e-01, 9.9945e-04, 1.5564e-03,  ..., 9.7275e-04,
           8.0490e-04, 2.9492e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4531e-01, 5.5420e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 1.2436e-03, 3.9062e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7031e-01, 5.9891e-04, 1.9360e-04,  ..., 1.5198e-02,
           0.0000e+00, 0.0000e+00],
          [5.8203e-01, 3.0518e-04, 1.7929e-04,  ..., 2.9907e-03,
           8.3618e-03, 0.0000e+00],
          [3.3789e-01, 1.6235e-02, 7.3242e-03,  ..., 5.1880e-03,
           4.3030e-03, 2.8906e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.6708e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.9591e-04, 2.0447e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7422e-01, 3.3760e-04, 1.4496e-04,  ..., 3.0975e-03,
           0.0000e+00, 0.0000e+00],
          [5.8594e-01, 2.6512e-04, 8.8692e-05,  ..., 1.8997e-03,
           3.3264e-03, 0.0000e+00],
          [3.9453e-01, 2.9297e-03, 1.6937e-03,  ..., 1.1673e-03,
           1.2054e-03, 2.9102e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.9302e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.6703e-03, 2.2430e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.0391e-01, 1.6235e-02, 5.2490e-03,  ..., 1.2894e-03,
           0.0000e+00, 0.0000e+00],
          [5.1172e-01, 3.1250e-02, 6.9580e-03,  ..., 8.4305e-04,
           2.7275e-04, 0.0000e+00],
          [3.9648e-01, 3.6774e-03, 3.7384e-03,  ..., 1.1215e-03,
           7.7057e-04, 2.8320e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.4556e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 4.9438e-03, 9.8267e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.4746e-01, 6.4087e-03, 9.3079e-04,  ..., 2.9419e-02,
           0.0000e+00, 0.0000e+00],
          [2.2656e-01, 8.3618e-03, 9.5749e-04,  ..., 2.0386e-02,
           1.6724e-02, 0.0000e+00],
          [4.5703e-01, 1.0605e-03, 1.2360e-03,  ..., 3.6621e-04,
           4.8447e-04, 2.6367e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.2114e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 2.9602e-03, 6.7749e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.8594e-01, 2.0447e-03, 4.6349e-04,  ..., 2.3499e-03,
           0.0000e+00, 0.0000e+00],
          [5.8594e-01, 1.9531e-03, 6.3324e-04,  ..., 1.1139e-03,
           1.7471e-03, 0.0000e+00],
          [3.9648e-01, 6.4697e-03, 2.3346e-03,  ..., 4.0588e-03,
           3.4637e-03, 2.7930e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.6060e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 3.2471e-02, 9.5215e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.0547e-02, 1.5545e-04, 1.1981e-05,  ..., 6.2500e-02,
           0.0000e+00, 0.0000e+00],
          [4.7302e-03, 6.2585e-06, 2.4587e-06,  ..., 8.9844e-01,
           7.1777e-02, 0.0000e+00],
          [3.6133e-01, 6.8054e-03, 4.2114e-03,  ..., 7.1411e-03,
           7.1411e-03, 2.8516e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.3977e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6484e-01, 2.2125e-03, 3.1982e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.3516e-01, 1.3828e-04, 7.3910e-05,  ..., 1.0315e-02,
           0.0000e+00, 0.0000e+00],
          [5.3906e-01, 4.1723e-05, 4.5776e-05,  ..., 4.7302e-03,
           5.9814e-03, 0.0000e+00],
          [4.1992e-01, 1.2894e-03, 6.6757e-04,  ..., 5.1880e-04,
           5.3787e-04, 2.8320e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.6172e-01, 2.3633e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 4.7302e-03, 3.4180e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7422e-01, 6.8283e-04, 3.7766e-04,  ..., 1.3794e-02,
           0.0000e+00, 0.0000e+00],
          [5.7031e-01, 3.8528e-04, 4.5013e-04,  ..., 1.3489e-02,
           1.9653e-02, 0.0000e+00],
          [3.9062e-01, 4.5166e-03, 3.3417e-03,  ..., 3.8452e-03,
           2.8992e-03, 2.8125e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 2.3926e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 1.2817e-03, 1.7242e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.9375e-01, 1.9989e-03, 7.0572e-04,  ..., 5.1498e-04,
           0.0000e+00, 0.0000e+00],
          [5.9375e-01, 3.4637e-03, 8.9264e-04,  ..., 3.5477e-04,
           2.0885e-04, 0.0000e+00],
          [4.1602e-01, 5.0354e-03, 1.0529e-03,  ..., 9.6130e-04,
           4.2725e-04, 2.7734e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 1.7548e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 1.0986e-03, 1.2817e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7812e-01, 9.3384e-03, 1.5945e-03,  ..., 3.9368e-03,
           0.0000e+00, 0.0000e+00],
          [5.4688e-01, 1.1353e-02, 2.4872e-03,  ..., 3.0975e-03,
           1.8845e-03, 0.0000e+00],
          [3.8867e-01, 7.1106e-03, 4.1809e-03,  ..., 4.9744e-03,
           3.4180e-03, 2.6758e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 9.7656e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.4414e-03, 2.6367e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7812e-01, 7.0953e-04, 4.0436e-04,  ..., 3.2806e-03,
           0.0000e+00, 0.0000e+00],
          [5.7812e-01, 3.0518e-04, 3.7956e-04,  ..., 2.3193e-03,
           3.8300e-03, 0.0000e+00],
          [4.1211e-01, 3.2806e-03, 1.1139e-03,  ..., 6.7520e-04,
           5.2643e-04, 2.8320e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.2024e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 3.5858e-03, 2.4567e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.4316e-01, 6.8359e-03, 1.0376e-03,  ..., 1.1719e-01,
           0.0000e+00, 0.0000e+00],
          [2.6953e-01, 7.6599e-03, 3.0518e-03,  ..., 3.9062e-02,
           5.4688e-02, 0.0000e+00],
          [2.8711e-01, 7.2632e-03, 6.3171e-03,  ..., 7.5073e-03,
           8.9111e-03, 3.1445e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2188e-01, 7.6172e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 2.8381e-03, 3.9368e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.5859e-01, 4.7874e-04, 1.2493e-04,  ..., 1.7014e-03,
           0.0000e+00, 0.0000e+00],
          [5.7031e-01, 5.9891e-04, 9.4414e-05,  ..., 7.5531e-04,
           6.3705e-04, 0.0000e+00],
          [3.9258e-01, 5.9509e-03, 1.9989e-03,  ..., 1.6022e-03,
           6.9046e-04, 2.9102e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 6.5002e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 8.2779e-04, 8.2779e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.1953e-01, 3.1982e-02, 4.0894e-03,  ..., 8.1177e-03,
           0.0000e+00, 0.0000e+00],
          [5.4688e-01, 2.7954e-02, 1.1902e-02,  ..., 1.2360e-03,
           6.9580e-03, 0.0000e+00],
          [4.4336e-01, 3.8147e-03, 1.3657e-03,  ..., 1.4954e-03,
           1.4038e-03, 2.6172e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.2196e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 2.0142e-03, 1.2054e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.6914e-01, 8.3984e-02, 2.0264e-02,  ..., 9.3994e-03,
           0.0000e+00, 0.0000e+00],
          [2.7344e-01, 1.3867e-01, 2.1118e-02,  ..., 5.0659e-03,
           2.3499e-03, 0.0000e+00],
          [2.6562e-01, 2.6733e-02, 2.5757e-02,  ..., 1.0559e-02,
           1.0071e-02, 2.2461e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.8281e-01, 1.1768e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 1.3885e-03, 3.0396e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.8594e-01, 4.7112e-04, 3.8910e-04,  ..., 4.9744e-03,
           0.0000e+00, 0.0000e+00],
          [6.0938e-01, 2.8038e-04, 1.5450e-04,  ..., 1.6327e-03,
           1.9150e-03, 0.0000e+00],
          [4.0430e-01, 4.0283e-03, 1.5793e-03,  ..., 1.7853e-03,
           1.6785e-03, 2.8516e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 7.9346e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 3.6133e-02, 9.5215e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [8.6426e-02, 2.6172e-01, 8.4229e-03,  ..., 5.8899e-03,
           0.0000e+00, 0.0000e+00],
          [5.2734e-02, 2.5977e-01, 1.3977e-02,  ..., 1.2756e-02,
           1.0681e-02, 0.0000e+00],
          [2.1387e-01, 6.4453e-02, 7.8735e-03,  ..., 6.5613e-03,
           4.4861e-03, 2.5781e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 2.4048e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9453e-01, 2.7924e-03, 1.0352e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.8594e-01, 4.6349e-04, 5.1117e-04,  ..., 4.2114e-03,
           0.0000e+00, 0.0000e+00],
          [5.7422e-01, 7.1716e-04, 4.4823e-04,  ..., 3.3112e-03,
           1.6632e-03, 0.0000e+00],
          [4.1602e-01, 6.3782e-03, 2.1362e-03,  ..., 5.9509e-04,
           2.0504e-04, 2.8320e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [5.7422e-01, 4.2578e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 1.7242e-03, 3.4180e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7812e-01, 2.6894e-04, 6.8188e-05,  ..., 2.3193e-03,
           0.0000e+00, 0.0000e+00],
          [5.8594e-01, 1.4114e-04, 3.5763e-05,  ..., 1.5717e-03,
           1.6708e-03, 0.0000e+00],
          [4.0039e-01, 5.5237e-03, 6.8283e-04,  ..., 7.0190e-04,
           6.8283e-04, 2.8906e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 2.3071e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.4766e-01, 1.8539e-03, 1.5039e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.1562e-01, 2.0447e-03, 4.8218e-03,  ..., 5.0964e-03,
           0.0000e+00, 0.0000e+00],
          [5.1562e-01, 1.1292e-03, 5.6458e-03,  ..., 9.2316e-04,
           2.0142e-03, 0.0000e+00],
          [3.8672e-01, 1.7471e-03, 1.0910e-03,  ..., 6.4468e-04,
           3.7766e-04, 2.9883e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.3977e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 1.5564e-03, 4.0527e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7812e-01, 1.4687e-04, 6.7902e-04,  ..., 1.3062e-02,
           0.0000e+00, 0.0000e+00],
          [5.6641e-01, 1.0633e-04, 2.8038e-04,  ..., 1.0864e-02,
           2.0386e-02, 0.0000e+00],
          [3.8086e-01, 2.6703e-03, 2.5024e-03,  ..., 1.1139e-03,
           9.1934e-04, 3.0078e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.6082e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 6.9809e-04, 4.6539e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.9766e-01, 4.9133e-03, 1.4114e-03,  ..., 4.4861e-03,
           0.0000e+00, 0.0000e+00],
          [5.7812e-01, 4.2725e-03, 1.6479e-03,  ..., 3.3264e-03,
           6.5308e-03, 0.0000e+00],
          [4.4531e-01, 1.2131e-03, 7.8583e-04,  ..., 5.2261e-04,
           2.9755e-04, 2.7148e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3750e-01, 6.1523e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.6406e-01, 8.4229e-03, 3.2617e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.3906e-01, 8.8215e-05, 9.1553e-04,  ..., 5.3787e-04,
           0.0000e+00, 0.0000e+00],
          [5.4688e-01, 8.2970e-05, 5.4169e-04,  ..., 1.9264e-04,
           2.6321e-04, 0.0000e+00],
          [3.7109e-01, 3.0670e-03, 2.8381e-03,  ..., 3.8338e-04,
           2.9945e-04, 3.0273e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.9368e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 6.5231e-04, 4.1199e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7031e-01, 1.4019e-04, 2.0409e-04,  ..., 9.2697e-04,
           0.0000e+00, 0.0000e+00],
          [5.8594e-01, 1.9932e-04, 1.4114e-04,  ..., 7.8964e-04,
           7.1716e-04, 0.0000e+00],
          [4.1406e-01, 3.0365e-03, 1.3885e-03,  ..., 5.9891e-04,
           2.5749e-04, 2.8320e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 2.1729e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 1.4465e-02, 1.2146e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.6641e-01, 1.2054e-03, 5.3406e-04,  ..., 7.9956e-03,
           0.0000e+00, 0.0000e+00],
          [5.7031e-01, 6.6376e-04, 3.4523e-04,  ..., 4.8218e-03,
           6.7139e-03, 0.0000e+00],
          [3.4961e-01, 8.6060e-03, 5.4626e-03,  ..., 5.8899e-03,
           3.7537e-03, 2.6758e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 5.7373e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 2.2095e-02, 2.1851e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.8047e-01, 3.2043e-03, 7.2861e-04,  ..., 6.8054e-03,
           0.0000e+00, 0.0000e+00],
          [4.7070e-01, 2.0447e-03, 5.2643e-04,  ..., 7.5989e-03,
           1.1047e-02, 0.0000e+00],
          [3.1641e-01, 5.2261e-04, 5.2261e-04,  ..., 1.2970e-03,
           1.5106e-03, 3.9648e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9062e-01, 1.0938e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 9.0942e-03, 3.4424e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.0820e-01, 2.6321e-04, 1.0681e-04,  ..., 7.2754e-02,
           0.0000e+00, 0.0000e+00],
          [3.0859e-01, 3.7956e-04, 1.9073e-04,  ..., 1.4355e-01,
           4.5654e-02, 0.0000e+00],
          [1.4648e-01, 2.0630e-02, 9.5825e-03,  ..., 9.5825e-03,
           1.0376e-02, 5.7031e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 5.9204e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 7.5150e-04, 3.1281e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7422e-01, 2.7657e-04, 2.7657e-04,  ..., 1.7456e-02,
           0.0000e+00, 0.0000e+00],
          [5.7031e-01, 2.9564e-04, 2.2411e-04,  ..., 1.5625e-02,
           5.0964e-03, 0.0000e+00],
          [4.2578e-01, 9.8419e-04, 1.0147e-03,  ..., 6.5613e-04,
           5.4550e-04, 2.7734e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 7.3242e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 3.1281e-04, 1.0910e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.0547e-01, 5.9605e-05, 8.1539e-05,  ..., 2.6093e-03,
           0.0000e+00, 0.0000e+00],
          [6.2891e-01, 4.4346e-05, 4.6968e-05,  ..., 8.3542e-04,
           6.4850e-04, 0.0000e+00],
          [4.1992e-01, 1.2207e-03, 5.7983e-04,  ..., 3.8528e-04,
           1.9360e-04, 2.9102e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 4.4441e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 3.7384e-04, 2.7275e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.8594e-01, 2.4128e-04, 1.0395e-04,  ..., 3.5095e-04,
           0.0000e+00, 0.0000e+00],
          [5.8984e-01, 3.0136e-04, 8.1062e-05,  ..., 4.0054e-04,
           7.8678e-05, 0.0000e+00],
          [4.0430e-01, 1.8387e-03, 7.2098e-04,  ..., 4.6539e-04,
           1.9360e-04, 2.9688e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 2.0752e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7656e-01, 4.5776e-03, 1.7212e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.3203e-01, 5.7602e-04, 1.6708e-03,  ..., 1.3770e-01,
           0.0000e+00, 0.0000e+00],
          [1.6699e-01, 5.5313e-04, 1.6556e-03,  ..., 3.0078e-01,
           1.0693e-01, 0.0000e+00],
          [1.7871e-01, 7.3853e-03, 1.2207e-02,  ..., 3.1738e-02,
           4.9561e-02, 3.4961e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.6060e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 2.0142e-02, 8.4839e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9023e-01, 2.5177e-03, 7.0953e-04,  ..., 5.3467e-02,
           0.0000e+00, 0.0000e+00],
          [4.2188e-01, 1.2131e-03, 4.5967e-04,  ..., 2.2339e-02,
           6.3477e-03, 0.0000e+00],
          [3.1641e-01, 2.1362e-02, 6.0120e-03,  ..., 9.6436e-03,
           5.1575e-03, 2.6758e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.9062e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.5869e-03, 1.5503e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.5078e-01, 1.5106e-03, 4.7302e-03,  ..., 1.6235e-02,
           0.0000e+00, 0.0000e+00],
          [5.3906e-01, 1.6022e-03, 7.2937e-03,  ..., 1.0986e-02,
           3.5156e-02, 0.0000e+00],
          [4.0625e-01, 2.5482e-03, 1.8082e-03,  ..., 1.5945e-03,
           1.8616e-03, 2.8125e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 7.3242e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 1.6708e-03, 1.1826e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.6250e-01, 3.6478e-05, 1.6689e-05,  ..., 5.3787e-04,
           0.0000e+00, 0.0000e+00],
          [5.6641e-01, 1.4484e-05, 8.2254e-06,  ..., 6.1417e-04,
           1.2436e-03, 0.0000e+00],
          [4.1406e-01, 7.4387e-04, 8.4305e-04,  ..., 6.1798e-04,
           3.9864e-04, 2.9297e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 3.9978e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 2.8076e-03, 1.4893e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.7891e-01, 1.5488e-03, 1.4893e-02,  ..., 5.7861e-02,
           0.0000e+00, 0.0000e+00],
          [2.9492e-01, 1.5182e-03, 1.3306e-02,  ..., 5.0781e-02,
           1.9238e-01, 0.0000e+00],
          [2.5195e-01, 7.9346e-03, 1.0681e-02,  ..., 1.7944e-02,
           2.6245e-02, 4.0430e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 3.1494e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 8.1787e-03, 1.7578e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.6641e-01, 1.4114e-03, 1.0300e-03,  ..., 2.0752e-02,
           0.0000e+00, 0.0000e+00],
          [5.8203e-01, 1.4267e-03, 9.2316e-04,  ..., 8.8501e-03,
           6.7139e-03, 0.0000e+00],
          [4.0625e-01, 2.9144e-03, 1.0757e-03,  ..., 8.1253e-04,
           4.3488e-04, 2.8711e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 4.1748e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.9141e-01, 1.4038e-02, 2.9492e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9609e-01, 4.5013e-04, 1.6708e-03,  ..., 4.2236e-02,
           0.0000e+00, 0.0000e+00],
          [5.2734e-01, 2.4223e-04, 9.9182e-04,  ..., 9.4604e-03,
           8.3008e-03, 0.0000e+00],
          [3.6328e-01, 5.0049e-03, 7.4158e-03,  ..., 2.3346e-03,
           1.9302e-03, 2.7734e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 6.5613e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 8.3923e-04, 5.6458e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.1719e-01, 1.3161e-04, 2.0981e-04,  ..., 1.4019e-04,
           0.0000e+00, 0.0000e+00],
          [6.3672e-01, 8.3447e-05, 1.0395e-04,  ..., 6.9141e-05,
           7.8678e-05, 0.0000e+00],
          [4.3359e-01, 6.4468e-04, 6.2561e-04,  ..., 5.7220e-04,
           4.5776e-04, 2.7148e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3359e-01, 6.6406e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.1094e-01, 1.3504e-03, 2.8906e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.5078e-01, 4.6349e-04, 1.5717e-03,  ..., 6.7139e-03,
           0.0000e+00, 0.0000e+00],
          [5.6641e-01, 2.6894e-04, 1.8082e-03,  ..., 1.0300e-03,
           2.7618e-03, 0.0000e+00],
          [3.9844e-01, 1.6327e-03, 8.2397e-04,  ..., 1.0777e-04,
           7.2002e-05, 2.9688e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.3640e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.1902e-03, 1.1475e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.2344e-01, 8.5831e-04, 6.6528e-03,  ..., 7.2021e-03,
           0.0000e+00, 0.0000e+00],
          [5.5469e-01, 4.2725e-04, 2.7313e-03,  ..., 3.0975e-03,
           2.7771e-03, 0.0000e+00],
          [1.4258e-01, 3.6774e-03, 6.4392e-03,  ..., 5.7678e-03,
           1.7212e-02, 5.2734e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2969e-01, 6.9824e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5312e-01, 4.1992e-02, 3.4943e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.9414e-01, 2.2430e-03, 1.8406e-04,  ..., 2.0386e-02,
           0.0000e+00, 0.0000e+00],
          [4.7070e-01, 2.7618e-03, 1.9932e-04,  ..., 2.8442e-02,
           1.7090e-02, 0.0000e+00],
          [3.9258e-01, 2.1484e-02, 2.6550e-03,  ..., 2.1973e-03,
           1.5106e-03, 2.7539e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 3.0708e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 4.7684e-04, 1.0757e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.1328e-01, 5.7459e-05, 4.3392e-05,  ..., 2.9182e-04,
           0.0000e+00, 0.0000e+00],
          [6.2500e-01, 3.6478e-05, 4.3869e-05,  ..., 1.5354e-04,
           7.0095e-05, 0.0000e+00],
          [4.5898e-01, 8.6975e-04, 9.8419e-04,  ..., 5.6076e-04,
           3.2997e-04, 2.6172e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 4.8218e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6875e-01, 1.4221e-02, 1.8433e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.1211e-01, 2.2095e-02, 7.7209e-03,  ..., 2.0386e-02,
           0.0000e+00, 0.0000e+00],
          [4.1406e-01, 2.3804e-02, 8.9111e-03,  ..., 1.3000e-02,
           1.6846e-02, 0.0000e+00],
          [4.1797e-01, 3.2043e-03, 1.7166e-03,  ..., 9.7656e-04,
           8.8882e-04, 2.6953e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4531e-01, 5.5176e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9844e-01, 9.7168e-02, 3.4027e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.4297e-01, 1.0872e-04, 2.8163e-06,  ..., 2.5940e-03,
           0.0000e+00, 0.0000e+00],
          [5.2734e-01, 3.0994e-05, 2.7120e-06,  ..., 1.4832e-02,
           8.1787e-03, 0.0000e+00],
          [4.1602e-01, 2.8534e-03, 1.1520e-03,  ..., 9.0027e-04,
           7.2479e-04, 2.8516e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 9.8877e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4922e-01, 2.0996e-02, 3.0029e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.5547e-01, 4.0627e-04, 3.0708e-04,  ..., 1.0059e-01,
           0.0000e+00, 0.0000e+00],
          [2.5391e-01, 2.4605e-04, 2.7084e-04,  ..., 1.0547e-01,
           1.8457e-01, 0.0000e+00],
          [4.7266e-01, 5.9891e-04, 2.6512e-04,  ..., 4.3678e-04,
           2.5749e-04, 2.5586e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.3306e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.5703e-01, 1.5991e-02, 2.6978e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7422e-01, 9.4986e-04, 4.3488e-04,  ..., 2.9755e-03,
           0.0000e+00, 0.0000e+00],
          [5.8594e-01, 1.6022e-03, 5.0354e-04,  ..., 1.5717e-03,
           8.2779e-04, 0.0000e+00],
          [3.9844e-01, 1.6327e-03, 4.8447e-04,  ..., 9.6130e-04,
           6.8283e-04, 3.0078e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8828e-01, 1.1597e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9609e-01, 1.5259e-03, 5.4169e-04,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.2109e-01, 4.7207e-05, 1.0550e-05,  ..., 1.2875e-04,
           0.0000e+00, 0.0000e+00],
          [6.0547e-01, 4.8876e-05, 2.0385e-05,  ..., 1.2493e-04,
           8.8692e-05, 0.0000e+00],
          [4.0039e-01, 1.0605e-03, 2.5940e-04,  ..., 4.9591e-05,
           9.8705e-05, 3.2422e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.0000e+00, 6.7139e-04, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 7.7438e-04, 1.3733e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [6.1719e-01, 1.1158e-04, 2.1667e-03,  ..., 9.9182e-04,
           0.0000e+00, 0.0000e+00],
          [6.2500e-01, 7.8678e-05, 1.5259e-03,  ..., 5.4550e-04,
           7.2098e-04, 0.0000e+00],
          [4.5312e-01, 1.9531e-03, 2.9297e-03,  ..., 3.8338e-04,
           2.6321e-04, 2.5391e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [6.7969e-01, 3.2227e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.8281e-01, 6.4453e-02, 5.1514e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [1.8848e-01, 5.3406e-04, 9.4223e-04,  ..., 3.2617e-01,
           0.0000e+00, 0.0000e+00],
          [5.1025e-02, 7.4387e-05, 2.7657e-04,  ..., 4.8438e-01,
           2.5195e-01, 0.0000e+00],
          [1.0254e-01, 6.5613e-03, 2.5757e-02,  ..., 3.5645e-02,
           4.8096e-02, 3.4570e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 6.6833e-03, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.7266e-01, 1.5793e-03, 2.6367e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7812e-01, 4.9210e-04, 2.3174e-04,  ..., 3.2501e-03,
           0.0000e+00, 0.0000e+00],
          [5.9766e-01, 3.1090e-04, 2.4128e-04,  ..., 1.2283e-03,
           9.5367e-04, 0.0000e+00],
          [4.7461e-01, 2.7161e-03, 5.6839e-04,  ..., 3.4523e-04,
           1.9741e-04, 2.4414e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.3047e-01, 2.7148e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.9219e-01, 3.0670e-03, 5.0354e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.7812e-01, 1.6251e-03, 6.0654e-04,  ..., 3.4943e-03,
           0.0000e+00, 0.0000e+00],
          [6.0547e-01, 1.1978e-03, 3.7575e-04,  ..., 6.8283e-04,
           7.0190e-04, 0.0000e+00],
          [4.1211e-01, 2.5940e-03, 4.8065e-04,  ..., 3.8719e-04,
           2.8229e-04, 2.8125e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6484e-01, 3.5889e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.6094e-01, 5.7678e-03, 3.3691e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.7891e-01, 2.5635e-03, 2.8992e-03,  ..., 2.7466e-02,
           0.0000e+00, 0.0000e+00],
          [3.6133e-01, 4.8828e-03, 4.8218e-03,  ..., 4.2480e-02,
           3.1738e-02, 0.0000e+00],
          [2.9785e-02, 1.0223e-03, 6.0791e-02,  ..., 4.0771e-02,
           3.5889e-02, 2.7344e-01]]]], device='cuda:0', dtype=torch.bfloat16), tensor([[[[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.8677e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2969e-01, 6.2988e-02, 7.1106e-03,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.6953e-01, 7.4768e-04, 7.2479e-04,  ..., 2.7344e-02,
           0.0000e+00, 0.0000e+00],
          [2.6367e-01, 8.2397e-04, 1.3199e-03,  ..., 4.2725e-02,
           4.4678e-02, 0.0000e+00],
          [1.1426e-01, 1.4771e-02, 1.5869e-02,  ..., 8.3496e-02,
           1.2598e-01, 1.6504e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.2188e-01, 7.7148e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.3359e-01, 4.7302e-03, 6.0791e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [5.8594e-01, 3.2616e-04, 3.8147e-04,  ..., 4.2419e-03,
           0.0000e+00, 0.0000e+00],
          [6.1328e-01, 3.4714e-04, 3.0708e-04,  ..., 7.0801e-03,
           8.4229e-03, 0.0000e+00],
          [3.5547e-01, 2.8076e-03, 2.2461e-02,  ..., 1.7548e-03,
           2.6398e-03, 2.9688e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 1.8799e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 5.5847e-03, 1.2573e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [3.2617e-01, 8.7891e-03, 8.7280e-03,  ..., 9.9121e-02,
           0.0000e+00, 0.0000e+00],
          [3.2031e-01, 8.1177e-03, 8.1787e-03,  ..., 9.1309e-02,
           8.6426e-02, 0.0000e+00],
          [1.5234e-01, 1.4893e-02, 2.6367e-02,  ..., 8.2031e-02,
           7.5195e-02, 1.4648e-01]],

         ...,

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [8.9062e-01, 1.0742e-01, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [7.3828e-01, 1.2793e-01, 1.3574e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.5898e-01, 2.1851e-02, 1.5869e-02,  ..., 1.2207e-02,
           0.0000e+00, 0.0000e+00],
          [4.5312e-01, 2.1362e-02, 1.4465e-02,  ..., 8.0566e-03,
           5.7983e-03, 0.0000e+00],
          [2.9492e-01, 1.8066e-02, 1.4038e-02,  ..., 6.5308e-03,
           7.6294e-03, 2.0215e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8438e-01, 1.6968e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [5.6641e-01, 1.3916e-02, 4.1992e-01,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [4.0430e-01, 1.1292e-02, 6.4392e-03,  ..., 1.3977e-02,
           0.0000e+00, 0.0000e+00],
          [3.7891e-01, 1.3733e-02, 8.6060e-03,  ..., 1.2268e-02,
           2.0996e-02, 0.0000e+00],
          [3.3789e-01, 2.9297e-02, 1.2207e-02,  ..., 2.8992e-03,
           2.0599e-03, 2.4121e-01]],

         [[1.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.4531e-01, 5.3467e-02, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [9.8047e-01, 4.4250e-03, 1.4709e-02,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.3535e-01, 1.0395e-04, 9.4414e-05,  ..., 8.6426e-02,
           0.0000e+00, 0.0000e+00],
          [1.5137e-01, 7.9632e-05, 5.6505e-05,  ..., 2.2656e-01,
           1.2109e-01, 0.0000e+00],
          [1.6992e-01, 2.8809e-02, 1.2573e-02,  ..., 9.7656e-03,
           1.8555e-02, 2.4414e-01]]]], device='cuda:0', dtype=torch.bfloat16))
<s> ▁Donald ▁was ▁very ▁good ▁at ▁playing ▁the ▁viol in ▁but ▁Matthew ▁was ▁not . ▁Matthew ▁gave ▁a ▁st unning ▁concert ▁performance .