test_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from molformer.utils import get_argparse_defaults, read_config\n",
    "from molformer.tokenizer import MolTranBertTokenizer\n",
    "from molformer.model.attention import TestBert\n",
    "from molformer.model.args import get_parser as ARGS\n",
    "from molformer.analysis import get_full_attention\n",
    "from attributedict.collections import AttributeDict\n",
    "import torch \n",
    "import pandas as pd\n",
    "import glob\n",
    "\n",
    "\n",
    "def get_bert(config, tokenizer):\n",
    "    bert_model = (\n",
    "        TestBert(\n",
    "            tokenizer.vocab,\n",
    "            config.seed_path,\n",
    "            rotate=config.rotate,\n",
    "            device=config.device,\n",
    "        )\n",
    "        .to(config.device)\n",
    "        .eval()\n",
    "    )\n",
    "    bert_model.load_state_dict(dict(torch.load(config.seed_path, map_location=config.device)[\"state_dict\"]), strict=True )\n",
    "    return bert_model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n"
     ]
    }
   ],
   "source": [
    "\n",
    "### defaults\n",
    "DEFAULTS = AttributeDict(get_argparse_defaults(ARGS()))\n",
    "TOKENIZER = MolTranBertTokenizer(\"finetune/bert_vocab.txt\")\n",
    "DEFAULTS.device = 'cpu'\n",
    "DEFAULTS.seed_path = '/Users/kevinmaikjablonka/Downloads/Pretrained MoLFormer/checkpoints/N-Step-Checkpoint_0_0.ckpt' # Doing this is kind of ugly. In the medium to long run this should be a real configuration system (e.g. Hydra)\n",
    "### import|model\n",
    "# register_attention(\"linear\")\n",
    "model = get_bert(DEFAULTS, TOKENIZER)\n",
    "\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized = TOKENIZER(\"CCO\", return_tensors=\"pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 4, 4, 4, 4, 1]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TOKENIZER('CCCC', return_tensors='pt')['input_ids']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 6.0439e-01, -7.6581e-01, -1.4492e+00,  1.3726e+00, -3.6720e-01,\n",
       "           4.9189e-01,  2.3665e+00,  1.4861e+00,  5.0351e-01, -2.0141e-01,\n",
       "          -1.9013e-01, -1.2264e+00, -5.7553e-01, -6.1704e-01,  1.9227e+00,\n",
       "           2.4361e-01,  2.3139e-01, -1.2067e+00,  1.6775e+00,  1.7688e+00,\n",
       "          -1.7363e-01, -9.3185e-01, -1.1856e-01,  2.8693e-01,  5.3709e-01,\n",
       "           1.2769e+00, -8.1280e-01, -7.5170e-01, -8.5207e-01, -6.0832e-01,\n",
       "           8.7988e-01,  6.1699e-01, -9.5774e-02,  7.9694e-01,  5.1645e-01,\n",
       "           9.5394e-01, -1.2341e-01,  1.1227e-01, -2.0075e+00,  3.8914e-01,\n",
       "          -1.5394e-01, -1.6298e+00, -1.3373e+00, -6.8660e-01,  3.0577e-02,\n",
       "           1.1332e+00,  1.7136e+00,  2.2455e-01, -2.4202e-01, -2.2155e+00,\n",
       "           1.5627e+00, -3.0705e-01, -1.2216e+00, -7.4428e-01,  2.9007e-01,\n",
       "           5.0757e-01, -1.9408e+00,  5.3317e-01,  1.5793e-01, -6.4825e-01,\n",
       "          -8.6230e-01,  2.7096e-01,  2.0881e-01,  5.6846e-01,  8.6085e-01,\n",
       "           7.8563e-01,  1.9426e+00, -2.0883e+00, -7.0042e-01, -4.9278e-01,\n",
       "          -3.7160e-01, -1.1541e+00, -9.7580e-01, -1.2484e+00, -2.0518e+00,\n",
       "           4.1033e-01, -1.8865e-01,  2.1873e-01,  1.0473e+00,  5.9461e-01,\n",
       "           3.7386e-02,  7.4167e-01, -1.3563e+00, -5.4585e-01,  8.2711e-01,\n",
       "           2.8196e-01, -1.0012e-01, -6.4009e-01,  4.1384e-02,  3.3165e-02,\n",
       "          -9.7335e-01, -8.9196e-01,  1.7649e-01,  9.6546e-02,  4.8163e-01,\n",
       "          -1.6893e+00,  3.9085e-01,  3.8173e-01, -3.8837e-01,  7.3510e-01,\n",
       "           3.0073e-01,  6.5975e-01, -6.1106e-01, -1.0344e+00, -3.1252e-01,\n",
       "          -6.3946e-01,  5.9014e-01, -1.6534e+00,  9.0243e-01, -3.4360e-01,\n",
       "          -1.3820e-01, -5.6923e-01, -2.9950e-01,  2.6911e-01, -4.6744e-01,\n",
       "           5.7265e-01, -6.6185e-02,  8.7908e-01, -2.2936e+00, -1.2854e+00,\n",
       "          -1.2655e+00,  3.0505e-01,  1.5737e+00, -1.7152e+00,  4.0742e-01,\n",
       "          -1.1600e+00, -6.5257e-01, -1.2059e+00, -3.1701e-02,  4.8480e-02,\n",
       "           7.4305e-01,  4.1721e-01, -1.3861e-01,  1.5755e-01,  5.2200e-01,\n",
       "           1.0137e+00,  1.5882e+00, -5.2017e-01, -1.4979e+00,  1.9182e+00,\n",
       "          -2.9210e-01,  1.5098e-02, -4.1774e-01, -9.3314e-01,  1.5611e+00,\n",
       "          -1.3891e+00,  6.1897e-02, -9.3792e-01,  1.0955e+00, -1.6320e+00,\n",
       "           1.4865e+00,  1.5980e+00,  1.2766e+00,  2.9774e-01, -6.4091e-01,\n",
       "           6.9366e-01, -2.7956e-01, -1.7125e+00, -5.3398e-02, -1.1431e+00,\n",
       "          -1.2136e+00,  1.6579e-01,  1.3312e-01, -1.0158e-01, -4.0860e-01,\n",
       "           2.9015e-01, -1.5957e+00,  3.3986e-01, -3.4227e-01,  1.2797e+00,\n",
       "          -1.4408e+00,  2.3856e+00,  5.7757e-01,  6.4686e-01,  1.4262e+00,\n",
       "          -2.3173e+00, -1.8978e+00, -6.0948e-01, -1.0771e+00,  1.2509e+00,\n",
       "          -1.1356e+00,  1.6598e+00, -8.7733e-01,  3.4966e-01,  1.4935e-01,\n",
       "           8.5350e-01,  5.5823e-01, -6.9673e-01, -1.0306e+00,  3.1953e-01,\n",
       "           4.9059e-02,  1.3867e-01,  1.4263e+00,  6.7125e-01,  5.1660e-01,\n",
       "           1.9159e+00, -9.9156e-01,  3.8680e-02, -8.7598e-01, -5.4850e-01,\n",
       "           1.4498e+00, -2.6673e-01,  2.3111e+00, -3.7443e-01, -3.8356e-01,\n",
       "          -5.0958e-01, -8.3753e-03, -1.3738e+00,  1.0877e+00, -3.2342e-01,\n",
       "          -1.9441e+00,  1.3446e+00, -1.1008e+00,  3.0784e-01, -2.5602e-01,\n",
       "          -8.6818e-01, -1.2530e+00,  1.4073e+00,  8.7100e-01,  1.0882e+00,\n",
       "          -1.3425e+00,  5.7286e-01,  7.1593e-01, -6.0662e-01,  3.3961e-01,\n",
       "           6.0332e-01, -5.7669e-01, -2.4630e+00,  1.0021e+00,  1.6365e-01,\n",
       "           3.6496e-01, -2.1101e+00,  3.9907e-01, -1.6686e+00,  1.2334e+00,\n",
       "           5.1509e-01, -6.2509e-01,  3.7880e-01,  7.8226e-01, -1.6364e+00,\n",
       "          -6.5811e-01, -1.3100e-01,  1.0778e+00,  1.0430e+00,  1.6198e+00,\n",
       "           1.4229e+00,  1.1558e+00,  2.5194e-01, -7.1475e-01,  9.2650e-01,\n",
       "          -5.4783e-01, -2.3959e+00, -9.6175e-01, -1.8424e+00, -3.1216e-01,\n",
       "           1.3171e+00,  1.9583e-01,  7.1182e-01, -1.0421e+00,  8.7440e-01,\n",
       "          -2.1606e-01, -1.5335e+00,  4.5983e-01, -5.8378e-01,  8.7261e-01,\n",
       "           6.4027e-01, -1.4275e+00,  3.8617e-01,  4.6730e-01,  8.4540e-03,\n",
       "           5.5940e-01,  4.7052e-01, -1.6799e-01,  5.1357e-01, -5.9003e-01,\n",
       "           1.1489e-01, -7.3467e-01,  8.3781e-01, -4.3062e-01,  2.1739e-01,\n",
       "           8.6634e-01, -2.7853e-01,  2.5282e-01,  1.9421e-01,  4.1622e-01,\n",
       "          -1.0051e+00,  3.7507e-01,  8.4164e-01, -1.7785e+00,  1.5524e-01,\n",
       "           1.4650e-02,  1.4641e+00,  3.1872e+00,  1.1748e+00,  7.1073e-01,\n",
       "           1.6377e-01, -1.5175e+00,  1.3455e+00, -1.9306e-01, -1.1851e+00,\n",
       "          -1.2459e+00, -2.7449e-01,  1.0321e-01,  2.0815e-01,  4.7443e-01,\n",
       "          -7.0130e-03, -1.4197e+00,  1.0839e+00, -2.0658e-01, -1.3663e+00,\n",
       "          -5.7764e-01, -8.7612e-01,  2.2769e+00,  5.1576e-02, -1.8398e-01,\n",
       "           1.9642e+00,  1.1415e+00,  7.7257e-01,  1.3726e+00,  5.6160e-01,\n",
       "           1.0682e+00, -9.4495e-01, -6.4916e-01, -1.2345e+00,  1.0927e-01,\n",
       "          -9.5403e-01, -1.4929e+00, -7.1732e-01, -8.9738e-01,  6.9039e-01,\n",
       "          -5.9452e-01,  2.1476e+00, -3.8512e-01,  1.3234e+00, -4.7145e-01,\n",
       "           1.1571e+00, -1.6474e-01, -7.6362e-01, -1.5938e+00,  5.1605e-01,\n",
       "           1.5171e+00, -6.0929e-02, -1.5410e+00,  1.0229e+00,  2.8818e-01,\n",
       "          -1.2548e+00,  9.5643e-01, -3.7722e-01, -1.1441e+00,  3.2453e-01,\n",
       "          -1.6318e+00, -5.9182e-01,  6.7567e-02,  4.8606e-01, -1.2274e+00,\n",
       "          -1.6232e-01,  1.3024e-01, -1.6313e-01,  1.4637e+00, -5.2927e-01,\n",
       "          -1.2536e-01, -9.2302e-02, -2.2909e-01, -7.0585e-01,  1.5419e+00,\n",
       "           8.0183e-01,  9.1345e-01,  3.1581e-01, -1.2555e+00,  1.3215e+00,\n",
       "          -2.1912e-02,  7.6191e-01, -3.9524e-01, -3.8774e-01,  1.5663e+00,\n",
       "           5.8161e-01, -5.1215e-01,  1.4901e+00, -1.3551e+00, -1.2865e+00,\n",
       "           1.1115e+00,  2.0003e-01, -1.7642e+00, -6.2223e-01, -7.6582e-01,\n",
       "          -7.8422e-01, -8.0116e-03, -1.5334e+00, -1.4540e+00,  2.6429e-01,\n",
       "          -1.9075e+00,  1.0778e-01,  7.3844e-01,  4.4241e-02,  1.3222e+00,\n",
       "           8.5433e-01,  1.4882e+00, -7.6827e-02,  8.9867e-01, -1.4134e+00,\n",
       "           2.0505e+00,  2.9952e-01, -6.3991e-01,  8.2976e-01,  6.7923e-01,\n",
       "           8.2050e-01, -2.0927e+00,  3.8254e-01, -9.6920e-01, -1.7914e-01,\n",
       "          -1.3767e+00, -2.4741e-01, -9.4061e-02,  7.0531e-01,  5.1540e-01,\n",
       "           1.3409e+00,  7.6676e-01,  5.6435e-01, -4.2486e-01,  3.2312e-01,\n",
       "           1.4450e+00, -1.7972e-01, -7.2643e-01,  1.0637e+00, -1.0995e+00,\n",
       "           8.2209e-01,  1.1855e+00,  7.8318e-01, -3.6015e-01,  7.6815e-01,\n",
       "          -9.5177e-01, -9.8592e-02,  5.3014e-01,  6.4492e-01,  2.0805e-01,\n",
       "           9.3617e-02, -1.3583e+00,  6.9968e-01,  1.1577e+00,  2.2353e+00,\n",
       "          -3.1763e-01, -1.2332e+00,  9.0994e-01,  3.8519e-01,  6.3384e-01,\n",
       "          -6.7771e-01, -2.7757e+00, -2.0151e+00,  7.2206e-01,  5.0434e-01,\n",
       "          -2.0941e+00,  2.8049e-01,  3.6703e-01, -7.7885e-02, -6.6514e-01,\n",
       "           1.0415e+00, -4.0465e-02,  1.0307e+00, -3.9897e-01, -5.3229e-01,\n",
       "           1.5099e+00, -3.7054e-01, -1.6660e-01,  1.2148e-01, -5.4900e-01,\n",
       "           5.1354e-01, -8.3725e-01,  2.1546e+00,  7.7091e-01, -1.6759e-01,\n",
       "          -3.0795e-01,  2.0143e-02, -2.1439e+00,  1.7369e+00, -1.8204e+00,\n",
       "          -2.3189e+00, -9.1580e-01, -3.6825e-01,  7.3360e-01, -5.2634e-01,\n",
       "           4.7852e-01, -1.6506e+00, -3.9879e-01,  6.6338e-01,  6.5531e-01,\n",
       "          -2.4629e+00,  8.4485e-01,  7.8110e-02,  8.6899e-01, -9.9955e-01,\n",
       "           6.4412e-01,  2.4216e-01, -1.5143e-01,  7.4650e-01,  1.1096e+00,\n",
       "           1.6302e+00, -1.5896e+00,  3.1406e-01,  2.2961e+00,  9.7439e-01,\n",
       "           1.5214e+00,  5.2758e-01,  1.3779e+00,  2.3096e-01,  3.6567e-02,\n",
       "           1.6863e+00, -1.6631e+00, -3.7338e-02,  1.1656e+00,  2.1073e-01,\n",
       "          -2.6450e+00, -4.0121e-01,  1.1805e+00,  6.4947e-02, -7.6755e-01,\n",
       "          -2.5010e-01,  6.5654e-01,  5.1902e-01,  2.3009e+00,  6.6910e-01,\n",
       "          -1.9800e+00,  8.2849e-01, -1.1805e+00,  5.4752e-01,  5.7545e-01,\n",
       "           9.3123e-01,  2.0781e-01,  9.6989e-01,  5.3822e-01, -2.7814e-01,\n",
       "          -1.4410e+00,  6.2056e-02,  4.8245e-01,  7.5104e-01, -1.2865e+00,\n",
       "          -8.9175e-01,  1.5663e-03, -5.8274e-01,  3.8205e-02,  6.1303e-01,\n",
       "           2.7409e-02,  5.6236e-01,  8.6177e-01, -8.2114e-01,  3.5944e-01,\n",
       "          -4.8478e-01, -1.1746e+00, -1.1479e+00, -8.0164e-02, -5.6227e-01,\n",
       "           7.5181e-01,  7.2413e-01, -1.3125e+00,  9.2079e-01,  1.5980e+00,\n",
       "           9.9285e-01,  3.6222e-01, -6.3108e-01,  2.8216e-01, -2.4054e+00,\n",
       "          -3.7496e-01, -1.8890e-01,  2.1146e-01,  6.4709e-01, -9.3073e-01,\n",
       "           1.3170e+00, -1.8674e-01,  1.4003e+00, -1.4113e+00, -3.2624e-01,\n",
       "           1.5865e+00,  1.1437e+00, -1.2791e+00,  1.2962e+00,  1.8923e+00,\n",
       "           6.9921e-01,  3.0989e-01,  9.2513e-01, -5.5923e-01,  2.2540e-01,\n",
       "           5.1967e-01, -4.7328e-01,  2.1545e+00,  3.0521e-01, -1.0574e+00,\n",
       "          -1.4904e+00,  1.0914e+00, -3.2095e-01,  1.2472e+00, -1.1885e+00,\n",
       "          -8.3387e-01,  6.6777e-01, -6.2476e-01,  7.1765e-01,  1.5651e-01,\n",
       "          -1.4138e+00, -2.3378e-01,  5.8180e-02,  1.3111e-01,  7.6946e-01,\n",
       "          -2.7545e-01,  1.0560e+00,  1.1419e+00,  2.6029e-01, -1.7187e+00,\n",
       "           9.4334e-01, -1.0189e-01,  5.3969e-01,  4.4992e-01, -1.6473e+00,\n",
       "          -3.0535e-01,  1.2289e+00,  8.8818e-01, -1.2626e+00, -1.7731e+00,\n",
       "           8.2798e-01,  6.6562e-01, -2.7702e-01,  2.0658e-01,  1.8366e+00,\n",
       "           1.1529e+00, -8.8290e-01,  2.8415e-01,  1.6937e-01,  3.0613e-01,\n",
       "          -7.9662e-01,  7.1522e-01, -4.9449e-01,  1.1401e+00,  1.8437e-01,\n",
       "           3.1806e-01,  1.8373e-01,  1.4082e+00,  1.4317e-01, -9.8794e-01,\n",
       "           8.3553e-01, -6.3228e-01,  5.8457e-01, -6.5340e-01, -1.5752e+00,\n",
       "           1.8460e-01,  7.0415e-02,  5.7021e-01, -2.5311e-01,  9.2941e-01,\n",
       "           7.6857e-01,  1.0223e+00,  2.1328e+00,  6.6590e-01,  2.1261e-01,\n",
       "           1.5542e+00, -5.8582e-02, -1.1036e+00, -3.1587e-01,  3.0971e-01,\n",
       "           1.1114e+00, -1.5064e+00,  5.1544e-01, -1.0841e+00,  4.8193e-01,\n",
       "          -1.1854e+00, -9.4728e-01,  5.1364e-01, -9.4556e-01, -8.5785e-01,\n",
       "          -2.9730e-02, -5.7436e-01,  2.3949e-01, -1.4107e+00, -1.5211e+00,\n",
       "           2.5066e-01, -1.5189e+00,  1.9048e-01,  4.3099e-01,  6.3987e-02,\n",
       "          -9.7191e-02,  1.1738e+00, -1.8237e-01, -2.3464e+00,  8.7769e-01,\n",
       "           1.0189e-01, -4.3297e-01,  2.1081e-01, -1.9473e+00, -7.3697e-02,\n",
       "           4.1624e-01, -3.9167e-01, -5.8175e-01, -8.8147e-01, -1.4844e-01,\n",
       "           1.2709e+00,  4.2625e-01,  1.8341e+00, -2.1334e+00, -1.2420e+00,\n",
       "          -8.9280e-01, -3.4863e-01, -3.5879e-01,  8.6531e-01, -1.5308e-01,\n",
       "           7.4393e-01,  6.3796e-01, -2.6612e-02,  6.2543e-01, -2.3934e-01,\n",
       "           5.8220e-01,  7.6726e-01, -2.6648e-01, -3.6464e-01, -1.9222e-01,\n",
       "           2.5277e+00, -2.6259e+00,  9.4998e-01,  1.9247e+00, -7.4660e-01,\n",
       "          -3.6201e-01,  7.8026e-01, -1.2349e-01, -1.4649e+00, -1.1008e+00,\n",
       "           1.1960e+00,  1.4072e+00, -9.9276e-01,  1.4210e-01,  6.9475e-01,\n",
       "          -2.2446e-01, -2.2372e-01, -1.2867e+00,  2.4577e-02,  8.5849e-01,\n",
       "           5.3319e-01,  2.5396e-01, -2.4462e+00, -8.5320e-01,  9.3896e-01,\n",
       "          -1.8952e+00, -7.0072e-01,  1.3124e+00, -1.0649e+00, -1.0026e-01,\n",
       "          -2.3362e-01, -1.6010e-01, -1.7674e-01, -6.2552e-01,  4.1223e-01,\n",
       "          -1.6723e+00, -9.1166e-01, -2.7198e-01, -8.5112e-01,  1.2085e+00,\n",
       "           3.8021e-01,  8.0832e-01,  3.5215e-02,  1.0752e-01, -9.3927e-01,\n",
       "           8.9779e-01, -3.0045e-01,  8.5201e-01, -8.9339e-01, -9.8239e-01,\n",
       "          -9.3384e-01, -6.9985e-01,  1.2777e+00, -2.9135e-02, -9.3991e-01,\n",
       "           2.9338e-02, -2.2081e-01, -4.4310e-01]], grad_fn=<SliceBackward0>),\n",
       " [tensor([[[[0.1691, 0.1537, 0.1693, 0.1897, 0.1949, 0.1234],\n",
       "            [0.0894, 0.1622, 0.1837, 0.2223, 0.2519, 0.0905],\n",
       "            [0.0908, 0.1679, 0.1843, 0.2222, 0.2537, 0.0811],\n",
       "            [0.1105, 0.1709, 0.1806, 0.2137, 0.2425, 0.0818],\n",
       "            [0.1278, 0.1676, 0.1753, 0.2027, 0.2247, 0.1019],\n",
       "            [0.3871, 0.1441, 0.1145, 0.1011, 0.0879, 0.1653]],\n",
       "  \n",
       "           [[0.0625, 0.2281, 0.2119, 0.1987, 0.1964, 0.1024],\n",
       "            [0.1443, 0.2088, 0.2020, 0.1786, 0.1577, 0.1086],\n",
       "            [0.1154, 0.1737, 0.1851, 0.1925, 0.1840, 0.1493],\n",
       "            [0.1036, 0.1480, 0.1562, 0.1876, 0.2089, 0.1957],\n",
       "            [0.0953, 0.1669, 0.1540, 0.1744, 0.2136, 0.1959],\n",
       "            [0.2232, 0.1604, 0.1372, 0.1420, 0.1759, 0.1611]],\n",
       "  \n",
       "           [[0.3681, 0.1322, 0.1402, 0.1408, 0.1357, 0.0830],\n",
       "            [0.1549, 0.1597, 0.1640, 0.1715, 0.1731, 0.1767],\n",
       "            [0.1250, 0.1781, 0.1773, 0.1779, 0.1771, 0.1646],\n",
       "            [0.1204, 0.1813, 0.1825, 0.1815, 0.1776, 0.1567],\n",
       "            [0.1376, 0.1684, 0.1772, 0.1829, 0.1799, 0.1539],\n",
       "            [0.0546, 0.1968, 0.1948, 0.1801, 0.1635, 0.2102]],\n",
       "  \n",
       "           [[0.1267, 0.2302, 0.1872, 0.1674, 0.1556, 0.1328],\n",
       "            [0.0857, 0.2147, 0.1852, 0.1899, 0.2171, 0.1074],\n",
       "            [0.0995, 0.2202, 0.1839, 0.1829, 0.2108, 0.1027],\n",
       "            [0.1092, 0.2233, 0.1876, 0.1829, 0.2061, 0.0909],\n",
       "            [0.1030, 0.2201, 0.1953, 0.1930, 0.2094, 0.0792],\n",
       "            [0.1797, 0.1688, 0.1404, 0.1469, 0.1707, 0.1935]],\n",
       "  \n",
       "           [[0.2654, 0.0590, 0.1152, 0.1422, 0.2028, 0.2154],\n",
       "            [0.1745, 0.0886, 0.1049, 0.1901, 0.2585, 0.1834],\n",
       "            [0.1664, 0.0951, 0.1081, 0.1932, 0.2656, 0.1716],\n",
       "            [0.1613, 0.1009, 0.1135, 0.1966, 0.2710, 0.1567],\n",
       "            [0.1594, 0.1083, 0.1227, 0.1995, 0.2697, 0.1404],\n",
       "            [0.0591, 0.1260, 0.1585, 0.2023, 0.2188, 0.2352]],\n",
       "  \n",
       "           [[0.0779, 0.1732, 0.1995, 0.2051, 0.2051, 0.1393],\n",
       "            [0.2156, 0.1236, 0.1425, 0.1418, 0.1500, 0.2264],\n",
       "            [0.2374, 0.1095, 0.1230, 0.1237, 0.1441, 0.2622],\n",
       "            [0.2984, 0.0869, 0.0973, 0.0990, 0.1242, 0.2943],\n",
       "            [0.3172, 0.0825, 0.0923, 0.0939, 0.1141, 0.3000],\n",
       "            [0.1885, 0.1611, 0.1785, 0.1908, 0.2009, 0.0802]],\n",
       "  \n",
       "           [[0.3203, 0.1265, 0.1155, 0.1121, 0.1149, 0.2109],\n",
       "            [0.1215, 0.1718, 0.1832, 0.1899, 0.1949, 0.1388],\n",
       "            [0.1128, 0.1683, 0.1810, 0.1876, 0.1925, 0.1578],\n",
       "            [0.1092, 0.1689, 0.1829, 0.1879, 0.1885, 0.1625],\n",
       "            [0.1122, 0.1683, 0.1837, 0.1876, 0.1836, 0.1646],\n",
       "            [0.2071, 0.1464, 0.1540, 0.1505, 0.1436, 0.1986]],\n",
       "  \n",
       "           [[0.0704, 0.1355, 0.1599, 0.2085, 0.2511, 0.1747],\n",
       "            [0.0480, 0.1928, 0.2037, 0.2219, 0.2470, 0.0866],\n",
       "            [0.0424, 0.1960, 0.2092, 0.2262, 0.2449, 0.0813],\n",
       "            [0.0428, 0.1957, 0.2131, 0.2305, 0.2436, 0.0743],\n",
       "            [0.0481, 0.1923, 0.2102, 0.2301, 0.2461, 0.0732],\n",
       "            [0.1198, 0.1634, 0.1842, 0.2014, 0.2043, 0.1268]],\n",
       "  \n",
       "           [[0.0511, 0.2125, 0.2152, 0.2179, 0.2192, 0.0840],\n",
       "            [0.1151, 0.2077, 0.1933, 0.1950, 0.2039, 0.0849],\n",
       "            [0.1166, 0.2206, 0.1960, 0.1907, 0.1949, 0.0811],\n",
       "            [0.1278, 0.2246, 0.1945, 0.1847, 0.1860, 0.0824],\n",
       "            [0.1327, 0.2245, 0.1941, 0.1832, 0.1848, 0.0806],\n",
       "            [0.1564, 0.2122, 0.1928, 0.1705, 0.1544, 0.1137]],\n",
       "  \n",
       "           [[0.0518, 0.2070, 0.2097, 0.2209, 0.2312, 0.0794],\n",
       "            [0.0910, 0.2233, 0.2239, 0.2234, 0.2207, 0.0177],\n",
       "            [0.0681, 0.2247, 0.2271, 0.2310, 0.2331, 0.0160],\n",
       "            [0.0724, 0.2242, 0.2277, 0.2317, 0.2336, 0.0103],\n",
       "            [0.0972, 0.2211, 0.2246, 0.2253, 0.2224, 0.0094],\n",
       "            [0.0908, 0.2086, 0.2167, 0.2183, 0.2189, 0.0467]],\n",
       "  \n",
       "           [[0.2218, 0.1096, 0.1100, 0.1081, 0.1003, 0.3502],\n",
       "            [0.3141, 0.0791, 0.1192, 0.1609, 0.1838, 0.1428],\n",
       "            [0.3249, 0.0757, 0.1137, 0.1532, 0.1732, 0.1594],\n",
       "            [0.3323, 0.0749, 0.1108, 0.1483, 0.1667, 0.1671],\n",
       "            [0.3444, 0.0757, 0.1092, 0.1448, 0.1623, 0.1636],\n",
       "            [0.2928, 0.1111, 0.1363, 0.1596, 0.1635, 0.1367]],\n",
       "  \n",
       "           [[0.1500, 0.1698, 0.1681, 0.1589, 0.1544, 0.1989],\n",
       "            [0.0479, 0.2187, 0.2221, 0.2163, 0.2039, 0.0912],\n",
       "            [0.0502, 0.2228, 0.2215, 0.2129, 0.2022, 0.0903],\n",
       "            [0.0597, 0.2157, 0.2200, 0.2127, 0.2027, 0.0892],\n",
       "            [0.0659, 0.2111, 0.2207, 0.2156, 0.2050, 0.0818],\n",
       "            [0.2065, 0.1816, 0.1850, 0.1769, 0.1665, 0.0835]]]]),\n",
       "  tensor([[[[0.2262, 0.1456, 0.1548, 0.1610, 0.1691, 0.1433],\n",
       "            [0.1635, 0.1570, 0.1656, 0.1792, 0.1841, 0.1507],\n",
       "            [0.1829, 0.1506, 0.1587, 0.1725, 0.1786, 0.1567],\n",
       "            [0.1996, 0.1489, 0.1589, 0.1753, 0.1843, 0.1330],\n",
       "            [0.2018, 0.1537, 0.1637, 0.1813, 0.1919, 0.1075],\n",
       "            [0.2542, 0.1595, 0.1525, 0.1450, 0.1366, 0.1522]],\n",
       "  \n",
       "           [[0.1412, 0.1917, 0.1800, 0.1522, 0.1579, 0.1771],\n",
       "            [0.0979, 0.1597, 0.1817, 0.1790, 0.1630, 0.2187],\n",
       "            [0.0917, 0.1607, 0.1781, 0.1758, 0.1606, 0.2331],\n",
       "            [0.0859, 0.1643, 0.1725, 0.1692, 0.1609, 0.2471],\n",
       "            [0.0837, 0.1679, 0.1659, 0.1602, 0.1608, 0.2613],\n",
       "            [0.1109, 0.1852, 0.2056, 0.1960, 0.1704, 0.1319]],\n",
       "  \n",
       "           [[0.1713, 0.1384, 0.1753, 0.1839, 0.1780, 0.1530],\n",
       "            [0.2196, 0.1166, 0.1408, 0.1474, 0.1388, 0.2368],\n",
       "            [0.1828, 0.1042, 0.1359, 0.1580, 0.1529, 0.2663],\n",
       "            [0.1538, 0.1291, 0.1388, 0.1493, 0.1500, 0.2791],\n",
       "            [0.1336, 0.1692, 0.1511, 0.1416, 0.1426, 0.2619],\n",
       "            [0.1499, 0.1624, 0.1999, 0.2112, 0.1832, 0.0934]],\n",
       "  \n",
       "           [[0.1472, 0.1350, 0.1547, 0.1628, 0.1611, 0.2393],\n",
       "            [0.1193, 0.1894, 0.1885, 0.1700, 0.1555, 0.1772],\n",
       "            [0.1450, 0.1513, 0.1646, 0.1634, 0.1623, 0.2133],\n",
       "            [0.1428, 0.1313, 0.1497, 0.1591, 0.1624, 0.2547],\n",
       "            [0.1180, 0.1347, 0.1533, 0.1627, 0.1635, 0.2680],\n",
       "            [0.2126, 0.0819, 0.1070, 0.1173, 0.1096, 0.3717]],\n",
       "  \n",
       "           [[0.0481, 0.2240, 0.2226, 0.1665, 0.1374, 0.2014],\n",
       "            [0.0703, 0.1955, 0.2036, 0.2125, 0.2132, 0.1049],\n",
       "            [0.0527, 0.1902, 0.2080, 0.2237, 0.2229, 0.1026],\n",
       "            [0.0523, 0.1922, 0.2062, 0.2183, 0.2170, 0.1140],\n",
       "            [0.0592, 0.1989, 0.1996, 0.2017, 0.2008, 0.1398],\n",
       "            [0.0678, 0.2215, 0.1981, 0.1875, 0.1980, 0.1271]],\n",
       "  \n",
       "           [[0.1934, 0.1462, 0.1629, 0.1557, 0.1415, 0.2002],\n",
       "            [0.0262, 0.1648, 0.2033, 0.2320, 0.2429, 0.1308],\n",
       "            [0.0397, 0.1651, 0.1981, 0.2233, 0.2469, 0.1268],\n",
       "            [0.0499, 0.1705, 0.1998, 0.2161, 0.2419, 0.1218],\n",
       "            [0.0547, 0.1768, 0.2048, 0.2114, 0.2298, 0.1226],\n",
       "            [0.2902, 0.1300, 0.1742, 0.1670, 0.1046, 0.1340]],\n",
       "  \n",
       "           [[0.1272, 0.1843, 0.1899, 0.1960, 0.2008, 0.1018],\n",
       "            [0.0910, 0.1772, 0.1865, 0.1908, 0.1909, 0.1636],\n",
       "            [0.0830, 0.1909, 0.1925, 0.1869, 0.1783, 0.1684],\n",
       "            [0.0769, 0.1974, 0.1985, 0.1908, 0.1792, 0.1573],\n",
       "            [0.0717, 0.1975, 0.2024, 0.1981, 0.1876, 0.1428],\n",
       "            [0.0412, 0.1912, 0.1981, 0.1860, 0.1664, 0.2170]],\n",
       "  \n",
       "           [[0.2433, 0.1519, 0.1850, 0.1742, 0.1063, 0.1392],\n",
       "            [0.1319, 0.2200, 0.2096, 0.1623, 0.1085, 0.1677],\n",
       "            [0.1230, 0.2051, 0.2208, 0.1911, 0.1100, 0.1499],\n",
       "            [0.1351, 0.1933, 0.2169, 0.1985, 0.1084, 0.1478],\n",
       "            [0.1551, 0.1943, 0.2021, 0.1906, 0.1195, 0.1385],\n",
       "            [0.2569, 0.1241, 0.1422, 0.1776, 0.1694, 0.1298]],\n",
       "  \n",
       "           [[0.1383, 0.2145, 0.2001, 0.1654, 0.1305, 0.1511],\n",
       "            [0.2214, 0.1280, 0.1453, 0.1590, 0.1586, 0.1878],\n",
       "            [0.2247, 0.1342, 0.1407, 0.1445, 0.1423, 0.2135],\n",
       "            [0.2096, 0.1600, 0.1517, 0.1314, 0.1106, 0.2368],\n",
       "            [0.1638, 0.1906, 0.1735, 0.1368, 0.1013, 0.2340],\n",
       "            [0.1397, 0.1506, 0.1615, 0.1617, 0.1429, 0.2436]],\n",
       "  \n",
       "           [[0.0706, 0.1785, 0.1896, 0.1983, 0.1945, 0.1686],\n",
       "            [0.0240, 0.1578, 0.1594, 0.1800, 0.2132, 0.2656],\n",
       "            [0.0203, 0.1588, 0.1587, 0.1730, 0.2055, 0.2837],\n",
       "            [0.0196, 0.1573, 0.1592, 0.1738, 0.2049, 0.2853],\n",
       "            [0.0143, 0.1553, 0.1578, 0.1763, 0.2107, 0.2856],\n",
       "            [0.1348, 0.1714, 0.1870, 0.1799, 0.1657, 0.1612]],\n",
       "  \n",
       "           [[0.1896, 0.0879, 0.1059, 0.1569, 0.1958, 0.2639],\n",
       "            [0.2702, 0.1193, 0.1374, 0.1628, 0.1855, 0.1248],\n",
       "            [0.2663, 0.1165, 0.1334, 0.1561, 0.1789, 0.1489],\n",
       "            [0.2848, 0.1074, 0.1224, 0.1429, 0.1668, 0.1757],\n",
       "            [0.3286, 0.0895, 0.1068, 0.1299, 0.1575, 0.1877],\n",
       "            [0.2506, 0.0995, 0.1202, 0.1382, 0.1658, 0.2257]],\n",
       "  \n",
       "           [[0.2641, 0.0999, 0.0944, 0.0959, 0.1092, 0.3366],\n",
       "            [0.1594, 0.1229, 0.1216, 0.1380, 0.1611, 0.2970],\n",
       "            [0.1504, 0.1371, 0.1181, 0.1372, 0.1657, 0.2915],\n",
       "            [0.1531, 0.1599, 0.1234, 0.1287, 0.1562, 0.2787],\n",
       "            [0.1611, 0.1831, 0.1316, 0.1179, 0.1402, 0.2661],\n",
       "            [0.4411, 0.0436, 0.0233, 0.0467, 0.1021, 0.3432]]]]),\n",
       "  tensor([[[[0.1531, 0.1475, 0.1673, 0.1896, 0.2016, 0.1409],\n",
       "            [0.2191, 0.1495, 0.1541, 0.1711, 0.1891, 0.1172],\n",
       "            [0.2081, 0.1644, 0.1635, 0.1750, 0.1891, 0.0998],\n",
       "            [0.1889, 0.1731, 0.1733, 0.1834, 0.1945, 0.0868],\n",
       "            [0.1677, 0.1705, 0.1797, 0.1953, 0.2070, 0.0799],\n",
       "            [0.3153, 0.1227, 0.1034, 0.1055, 0.1157, 0.2373]],\n",
       "  \n",
       "           [[0.1156, 0.1619, 0.1584, 0.1622, 0.1643, 0.2376],\n",
       "            [0.1860, 0.1190, 0.1063, 0.0952, 0.0895, 0.4040],\n",
       "            [0.1852, 0.1229, 0.1119, 0.1028, 0.0967, 0.3803],\n",
       "            [0.1779, 0.1278, 0.1170, 0.1104, 0.1063, 0.3607],\n",
       "            [0.1770, 0.1365, 0.1230, 0.1152, 0.1121, 0.3361],\n",
       "            [0.1550, 0.1270, 0.1169, 0.1284, 0.1516, 0.3211]],\n",
       "  \n",
       "           [[0.0928, 0.1621, 0.1540, 0.1518, 0.1486, 0.2908],\n",
       "            [0.2891, 0.2130, 0.1497, 0.1197, 0.0827, 0.1459],\n",
       "            [0.2815, 0.2050, 0.1523, 0.1266, 0.0876, 0.1471],\n",
       "            [0.2704, 0.1916, 0.1536, 0.1334, 0.0943, 0.1566],\n",
       "            [0.2442, 0.1774, 0.1579, 0.1446, 0.1087, 0.1671],\n",
       "            [0.1642, 0.1842, 0.2026, 0.1947, 0.1603, 0.0940]],\n",
       "  \n",
       "           [[0.0877, 0.2329, 0.2068, 0.1800, 0.1671, 0.1255],\n",
       "            [0.0981, 0.1379, 0.1522, 0.1315, 0.1037, 0.3766],\n",
       "            [0.1201, 0.1167, 0.1206, 0.1181, 0.1187, 0.4058],\n",
       "            [0.1269, 0.1108, 0.0980, 0.1026, 0.1288, 0.4331],\n",
       "            [0.1014, 0.1500, 0.1147, 0.1033, 0.1308, 0.3999],\n",
       "            [0.0374, 0.2264, 0.2128, 0.1832, 0.1647, 0.1755]],\n",
       "  \n",
       "           [[0.2415, 0.1354, 0.1443, 0.1523, 0.1484, 0.1782],\n",
       "            [0.2526, 0.1270, 0.1326, 0.1255, 0.1089, 0.2533],\n",
       "            [0.2572, 0.1245, 0.1274, 0.1199, 0.1035, 0.2676],\n",
       "            [0.2282, 0.1293, 0.1299, 0.1220, 0.1060, 0.2846],\n",
       "            [0.1873, 0.1359, 0.1382, 0.1315, 0.1161, 0.2911],\n",
       "            [0.1400, 0.1676, 0.1964, 0.2041, 0.1917, 0.1003]],\n",
       "  \n",
       "           [[0.1551, 0.1656, 0.1752, 0.1727, 0.1488, 0.1826],\n",
       "            [0.2414, 0.1040, 0.1180, 0.1229, 0.1279, 0.2857],\n",
       "            [0.2544, 0.0997, 0.1254, 0.1237, 0.1156, 0.2813],\n",
       "            [0.2512, 0.1005, 0.1280, 0.1272, 0.1167, 0.2764],\n",
       "            [0.2605, 0.1057, 0.1191, 0.1182, 0.1281, 0.2684],\n",
       "            [0.1947, 0.1537, 0.1585, 0.1450, 0.1539, 0.1942]],\n",
       "  \n",
       "           [[0.0835, 0.1664, 0.1917, 0.2243, 0.2469, 0.0872],\n",
       "            [0.0748, 0.2873, 0.2448, 0.1827, 0.1563, 0.0541],\n",
       "            [0.0737, 0.2844, 0.2450, 0.1847, 0.1582, 0.0541],\n",
       "            [0.0728, 0.2852, 0.2449, 0.1843, 0.1581, 0.0547],\n",
       "            [0.0696, 0.2946, 0.2460, 0.1783, 0.1530, 0.0585],\n",
       "            [0.0831, 0.1988, 0.2026, 0.2041, 0.2064, 0.1051]],\n",
       "  \n",
       "           [[0.1304, 0.2182, 0.2009, 0.1527, 0.1126, 0.1851],\n",
       "            [0.2106, 0.1671, 0.1560, 0.1231, 0.0808, 0.2624],\n",
       "            [0.2880, 0.1560, 0.1410, 0.1051, 0.0626, 0.2474],\n",
       "            [0.3691, 0.1419, 0.1245, 0.0900, 0.0496, 0.2250],\n",
       "            [0.4089, 0.1281, 0.1077, 0.0855, 0.0530, 0.2168],\n",
       "            [0.3121, 0.1313, 0.1479, 0.1284, 0.0822, 0.1980]],\n",
       "  \n",
       "           [[0.1434, 0.1577, 0.1539, 0.1456, 0.1370, 0.2625],\n",
       "            [0.1514, 0.1912, 0.1887, 0.1901, 0.1912, 0.0875],\n",
       "            [0.1435, 0.1916, 0.1915, 0.1928, 0.1928, 0.0878],\n",
       "            [0.1292, 0.1888, 0.1917, 0.1945, 0.1953, 0.1006],\n",
       "            [0.1222, 0.1877, 0.1911, 0.1930, 0.1931, 0.1129],\n",
       "            [0.1388, 0.1942, 0.1922, 0.1835, 0.1717, 0.1196]],\n",
       "  \n",
       "           [[0.3239, 0.1222, 0.1349, 0.1246, 0.1073, 0.1872],\n",
       "            [0.2036, 0.1644, 0.1650, 0.1502, 0.1388, 0.1780],\n",
       "            [0.1678, 0.1558, 0.1616, 0.1546, 0.1470, 0.2133],\n",
       "            [0.1675, 0.1545, 0.1589, 0.1525, 0.1463, 0.2203],\n",
       "            [0.2067, 0.1526, 0.1539, 0.1445, 0.1375, 0.2049],\n",
       "            [0.3126, 0.1834, 0.1693, 0.1378, 0.1170, 0.0799]],\n",
       "  \n",
       "           [[0.2553, 0.2576, 0.1626, 0.0702, 0.0847, 0.1696],\n",
       "            [0.2952, 0.1217, 0.1409, 0.1564, 0.1817, 0.1042],\n",
       "            [0.2823, 0.1335, 0.1469, 0.1543, 0.1750, 0.1080],\n",
       "            [0.2811, 0.1387, 0.1513, 0.1577, 0.1749, 0.0963],\n",
       "            [0.2864, 0.1376, 0.1522, 0.1637, 0.1801, 0.0800],\n",
       "            [0.2955, 0.1268, 0.1229, 0.1276, 0.1639, 0.1634]],\n",
       "  \n",
       "           [[0.1594, 0.1448, 0.1553, 0.1858, 0.2149, 0.1397],\n",
       "            [0.1791, 0.1282, 0.1609, 0.2084, 0.2334, 0.0901],\n",
       "            [0.1800, 0.1254, 0.1532, 0.2023, 0.2330, 0.1060],\n",
       "            [0.1705, 0.1290, 0.1520, 0.1976, 0.2271, 0.1238],\n",
       "            [0.1489, 0.1411, 0.1569, 0.1949, 0.2165, 0.1417],\n",
       "            [0.0594, 0.2368, 0.1897, 0.1702, 0.1781, 0.1657]]]]),\n",
       "  tensor([[[[0.1145, 0.1770, 0.1687, 0.1767, 0.2199, 0.1433],\n",
       "            [0.2233, 0.1775, 0.1516, 0.1325, 0.1546, 0.1605],\n",
       "            [0.2433, 0.1787, 0.1497, 0.1253, 0.1452, 0.1577],\n",
       "            [0.2739, 0.1719, 0.1445, 0.1173, 0.1352, 0.1572],\n",
       "            [0.2915, 0.1635, 0.1378, 0.1107, 0.1284, 0.1682],\n",
       "            [0.0391, 0.2506, 0.2123, 0.1909, 0.2189, 0.0883]],\n",
       "  \n",
       "           [[0.2090, 0.1628, 0.1717, 0.1853, 0.1885, 0.0826],\n",
       "            [0.2782, 0.1796, 0.1625, 0.1574, 0.1559, 0.0664],\n",
       "            [0.3045, 0.1813, 0.1665, 0.1541, 0.1463, 0.0473],\n",
       "            [0.2929, 0.1769, 0.1760, 0.1648, 0.1530, 0.0363],\n",
       "            [0.2449, 0.1731, 0.1862, 0.1834, 0.1729, 0.0395],\n",
       "            [0.1063, 0.1885, 0.2060, 0.2089, 0.2042, 0.0861]],\n",
       "  \n",
       "           [[0.0789, 0.1324, 0.1220, 0.1585, 0.2384, 0.2697],\n",
       "            [0.0684, 0.1892, 0.1370, 0.1160, 0.1672, 0.3223],\n",
       "            [0.0696, 0.1958, 0.1413, 0.1131, 0.1565, 0.3238],\n",
       "            [0.0714, 0.2068, 0.1507, 0.1120, 0.1433, 0.3157],\n",
       "            [0.0720, 0.2114, 0.1563, 0.1177, 0.1466, 0.2961],\n",
       "            [0.0398, 0.1772, 0.1606, 0.1823, 0.2430, 0.1969]],\n",
       "  \n",
       "           [[0.2753, 0.2119, 0.1564, 0.1229, 0.1284, 0.1051],\n",
       "            [0.2259, 0.0996, 0.0869, 0.0992, 0.1564, 0.3321],\n",
       "            [0.2349, 0.0892, 0.0746, 0.0902, 0.1585, 0.3527],\n",
       "            [0.2365, 0.0791, 0.0631, 0.0820, 0.1587, 0.3805],\n",
       "            [0.2304, 0.0739, 0.0536, 0.0741, 0.1576, 0.4103],\n",
       "            [0.1653, 0.1737, 0.1328, 0.1198, 0.1201, 0.2884]],\n",
       "  \n",
       "           [[0.3193, 0.0500, 0.0558, 0.0906, 0.0999, 0.3844],\n",
       "            [0.3011, 0.0857, 0.0783, 0.0805, 0.1072, 0.3472],\n",
       "            [0.2730, 0.0760, 0.0705, 0.0756, 0.1098, 0.3952],\n",
       "            [0.2179, 0.0966, 0.0904, 0.0935, 0.1205, 0.3811],\n",
       "            [0.1921, 0.1106, 0.0995, 0.0972, 0.1166, 0.3841],\n",
       "            [0.1703, 0.1531, 0.1555, 0.1633, 0.1763, 0.1814]],\n",
       "  \n",
       "           [[0.1097, 0.1729, 0.1753, 0.1801, 0.1684, 0.1937],\n",
       "            [0.1195, 0.1857, 0.1763, 0.1835, 0.2002, 0.1348],\n",
       "            [0.1224, 0.1897, 0.1778, 0.1830, 0.2020, 0.1252],\n",
       "            [0.1269, 0.2061, 0.1849, 0.1795, 0.1919, 0.1107],\n",
       "            [0.1343, 0.2292, 0.1962, 0.1730, 0.1693, 0.0980],\n",
       "            [0.1595, 0.1924, 0.1831, 0.1799, 0.1725, 0.1128]],\n",
       "  \n",
       "           [[0.3430, 0.0906, 0.1027, 0.1129, 0.1119, 0.2389],\n",
       "            [0.2412, 0.1518, 0.1369, 0.1342, 0.1568, 0.1792],\n",
       "            [0.2272, 0.1640, 0.1434, 0.1338, 0.1502, 0.1814],\n",
       "            [0.2440, 0.1639, 0.1397, 0.1275, 0.1409, 0.1839],\n",
       "            [0.2899, 0.1364, 0.1190, 0.1175, 0.1355, 0.2017],\n",
       "            [0.2849, 0.1049, 0.0952, 0.1076, 0.1387, 0.2688]],\n",
       "  \n",
       "           [[0.1810, 0.1593, 0.1466, 0.1529, 0.1677, 0.1925],\n",
       "            [0.1007, 0.1654, 0.1839, 0.2141, 0.2239, 0.1119],\n",
       "            [0.0879, 0.1644, 0.1863, 0.2204, 0.2336, 0.1074],\n",
       "            [0.0709, 0.1674, 0.1902, 0.2264, 0.2427, 0.1024],\n",
       "            [0.0519, 0.1729, 0.1959, 0.2334, 0.2522, 0.0939],\n",
       "            [0.1627, 0.1628, 0.1759, 0.2009, 0.2177, 0.0800]],\n",
       "  \n",
       "           [[0.1268, 0.2124, 0.1891, 0.1744, 0.1689, 0.1283],\n",
       "            [0.1788, 0.1658, 0.1621, 0.1576, 0.1500, 0.1857],\n",
       "            [0.1614, 0.1845, 0.1749, 0.1629, 0.1507, 0.1656],\n",
       "            [0.1408, 0.2182, 0.1940, 0.1671, 0.1476, 0.1322],\n",
       "            [0.1335, 0.2354, 0.2043, 0.1698, 0.1465, 0.1105],\n",
       "            [0.1165, 0.2245, 0.1956, 0.1740, 0.1622, 0.1271]],\n",
       "  \n",
       "           [[0.1737, 0.1572, 0.1756, 0.1917, 0.1872, 0.1146],\n",
       "            [0.1405, 0.1676, 0.1741, 0.1952, 0.2075, 0.1152],\n",
       "            [0.1328, 0.1583, 0.1654, 0.2018, 0.2364, 0.1053],\n",
       "            [0.1231, 0.1653, 0.1694, 0.2036, 0.2394, 0.0992],\n",
       "            [0.1121, 0.1842, 0.1803, 0.1987, 0.2211, 0.1035],\n",
       "            [0.1289, 0.1659, 0.1360, 0.1319, 0.1312, 0.3061]],\n",
       "  \n",
       "           [[0.0709, 0.1384, 0.1996, 0.2557, 0.2491, 0.0864],\n",
       "            [0.0691, 0.1563, 0.1723, 0.1923, 0.1882, 0.2218],\n",
       "            [0.0781, 0.1595, 0.1750, 0.1882, 0.1801, 0.2191],\n",
       "            [0.0889, 0.1623, 0.1793, 0.1874, 0.1759, 0.2063],\n",
       "            [0.0968, 0.1660, 0.1844, 0.1890, 0.1748, 0.1891],\n",
       "            [0.1156, 0.1753, 0.2021, 0.2055, 0.1835, 0.1179]],\n",
       "  \n",
       "           [[0.1180, 0.1301, 0.1719, 0.1914, 0.1661, 0.2224],\n",
       "            [0.0922, 0.1604, 0.1726, 0.1478, 0.0994, 0.3277],\n",
       "            [0.0763, 0.1328, 0.1648, 0.1720, 0.1435, 0.3105],\n",
       "            [0.0840, 0.1181, 0.1627, 0.1896, 0.1751, 0.2706],\n",
       "            [0.1011, 0.1155, 0.1630, 0.1909, 0.1793, 0.2503],\n",
       "            [0.1756, 0.1523, 0.1742, 0.1863, 0.1732, 0.1385]]]]),\n",
       "  tensor([[[[0.2561, 0.1558, 0.1601, 0.1383, 0.1460, 0.1437],\n",
       "            [0.2566, 0.1276, 0.1256, 0.1098, 0.1306, 0.2497],\n",
       "            [0.2355, 0.1657, 0.1457, 0.0987, 0.0945, 0.2599],\n",
       "            [0.2086, 0.1903, 0.1673, 0.1050, 0.0809, 0.2480],\n",
       "            [0.1972, 0.1906, 0.1751, 0.1142, 0.0841, 0.2387],\n",
       "            [0.2650, 0.2003, 0.1868, 0.1136, 0.0988, 0.1354]],\n",
       "  \n",
       "           [[0.2039, 0.1610, 0.1600, 0.1792, 0.1587, 0.1372],\n",
       "            [0.1768, 0.1463, 0.1510, 0.1749, 0.1783, 0.1727],\n",
       "            [0.1813, 0.1502, 0.1603, 0.1740, 0.1652, 0.1691],\n",
       "            [0.1755, 0.1533, 0.1707, 0.1806, 0.1602, 0.1597],\n",
       "            [0.1615, 0.1516, 0.1766, 0.1919, 0.1682, 0.1502],\n",
       "            [0.1807, 0.1483, 0.1646, 0.1961, 0.1654, 0.1449]],\n",
       "  \n",
       "           [[0.1186, 0.1771, 0.1839, 0.1923, 0.1915, 0.1366],\n",
       "            [0.2329, 0.1708, 0.1490, 0.1468, 0.1655, 0.1351],\n",
       "            [0.2262, 0.1697, 0.1575, 0.1521, 0.1616, 0.1330],\n",
       "            [0.2063, 0.1636, 0.1662, 0.1649, 0.1676, 0.1313],\n",
       "            [0.2028, 0.1651, 0.1676, 0.1680, 0.1730, 0.1235],\n",
       "            [0.1372, 0.2000, 0.1642, 0.1461, 0.1593, 0.1931]],\n",
       "  \n",
       "           [[0.0992, 0.2153, 0.2169, 0.2117, 0.2081, 0.0488],\n",
       "            [0.1347, 0.2213, 0.2131, 0.1986, 0.1959, 0.0364],\n",
       "            [0.1339, 0.2218, 0.2138, 0.1983, 0.1920, 0.0401],\n",
       "            [0.1355, 0.2222, 0.2137, 0.1975, 0.1890, 0.0421],\n",
       "            [0.1320, 0.2213, 0.2130, 0.1975, 0.1891, 0.0470],\n",
       "            [0.1928, 0.1946, 0.1762, 0.1550, 0.1492, 0.1322]],\n",
       "  \n",
       "           [[0.1189, 0.1963, 0.1701, 0.1480, 0.1380, 0.2286],\n",
       "            [0.0683, 0.1369, 0.1672, 0.1887, 0.1930, 0.2458],\n",
       "            [0.0595, 0.1255, 0.1600, 0.1887, 0.1994, 0.2668],\n",
       "            [0.0537, 0.1239, 0.1579, 0.1871, 0.1996, 0.2777],\n",
       "            [0.0530, 0.1309, 0.1617, 0.1860, 0.1952, 0.2731],\n",
       "            [0.1055, 0.2405, 0.2233, 0.1769, 0.1197, 0.1341]],\n",
       "  \n",
       "           [[0.2281, 0.1466, 0.1409, 0.1230, 0.1186, 0.2427],\n",
       "            [0.2115, 0.1670, 0.1693, 0.1500, 0.1194, 0.1828],\n",
       "            [0.1673, 0.1812, 0.1833, 0.1608, 0.1281, 0.1793],\n",
       "            [0.1391, 0.1951, 0.1998, 0.1788, 0.1461, 0.1412],\n",
       "            [0.1490, 0.1864, 0.1960, 0.1884, 0.1600, 0.1203],\n",
       "            [0.2551, 0.1918, 0.1717, 0.1593, 0.1455, 0.0766]],\n",
       "  \n",
       "           [[0.0813, 0.2877, 0.2463, 0.1649, 0.0892, 0.1306],\n",
       "            [0.1071, 0.2224, 0.1978, 0.1634, 0.1352, 0.1741],\n",
       "            [0.0717, 0.2289, 0.2179, 0.1989, 0.1815, 0.1012],\n",
       "            [0.0499, 0.2108, 0.2181, 0.2166, 0.2127, 0.0919],\n",
       "            [0.0441, 0.1805, 0.2187, 0.2313, 0.2261, 0.0993],\n",
       "            [0.0950, 0.2067, 0.1992, 0.1742, 0.1568, 0.1682]],\n",
       "  \n",
       "           [[0.1496, 0.1080, 0.1673, 0.1919, 0.1626, 0.2206],\n",
       "            [0.1563, 0.1221, 0.1864, 0.1991, 0.1375, 0.1985],\n",
       "            [0.1206, 0.0950, 0.1949, 0.2295, 0.1628, 0.1972],\n",
       "            [0.1335, 0.0812, 0.1744, 0.2195, 0.1721, 0.2192],\n",
       "            [0.1864, 0.0845, 0.1399, 0.1766, 0.1553, 0.2574],\n",
       "            [0.1669, 0.1248, 0.1950, 0.2341, 0.1963, 0.0829]],\n",
       "  \n",
       "           [[0.3273, 0.0885, 0.1278, 0.1620, 0.1971, 0.0973],\n",
       "            [0.3237, 0.1152, 0.1226, 0.1577, 0.2049, 0.0758],\n",
       "            [0.3178, 0.1171, 0.1255, 0.1523, 0.1893, 0.0979],\n",
       "            [0.3114, 0.1322, 0.1456, 0.1522, 0.1521, 0.1065],\n",
       "            [0.3165, 0.1413, 0.1614, 0.1555, 0.1290, 0.0963],\n",
       "            [0.2824, 0.1200, 0.1392, 0.1442, 0.2444, 0.0698]],\n",
       "  \n",
       "           [[0.1416, 0.1854, 0.2190, 0.2070, 0.1578, 0.0892],\n",
       "            [0.0312, 0.1746, 0.2059, 0.2007, 0.1995, 0.1881],\n",
       "            [0.0230, 0.1998, 0.2280, 0.2003, 0.1752, 0.1737],\n",
       "            [0.0555, 0.2323, 0.2597, 0.1865, 0.1102, 0.1557],\n",
       "            [0.0839, 0.2174, 0.2504, 0.1856, 0.1009, 0.1618],\n",
       "            [0.1594, 0.1597, 0.1904, 0.1762, 0.1935, 0.1207]],\n",
       "  \n",
       "           [[0.2553, 0.2084, 0.1520, 0.1276, 0.1400, 0.1168],\n",
       "            [0.0844, 0.2148, 0.1765, 0.1624, 0.1841, 0.1778],\n",
       "            [0.0718, 0.2200, 0.1779, 0.1601, 0.1751, 0.1952],\n",
       "            [0.0655, 0.2234, 0.1835, 0.1611, 0.1652, 0.2012],\n",
       "            [0.0732, 0.2201, 0.1891, 0.1650, 0.1568, 0.1958],\n",
       "            [0.0955, 0.1413, 0.1512, 0.1754, 0.1903, 0.2463]],\n",
       "  \n",
       "           [[0.2564, 0.0799, 0.1017, 0.1275, 0.1647, 0.2698],\n",
       "            [0.1676, 0.1117, 0.1250, 0.1670, 0.1990, 0.2298],\n",
       "            [0.1662, 0.1213, 0.1348, 0.1671, 0.1884, 0.2223],\n",
       "            [0.1687, 0.1296, 0.1415, 0.1660, 0.1800, 0.2142],\n",
       "            [0.1738, 0.1385, 0.1470, 0.1632, 0.1724, 0.2052],\n",
       "            [0.1447, 0.1899, 0.1939, 0.1876, 0.1651, 0.1190]]]]),\n",
       "  tensor([[[[0.2913, 0.1534, 0.1577, 0.1528, 0.1400, 0.1048],\n",
       "            [0.2851, 0.1437, 0.1312, 0.1218, 0.1192, 0.1990],\n",
       "            [0.2868, 0.1554, 0.1381, 0.1239, 0.1193, 0.1764],\n",
       "            [0.2569, 0.1722, 0.1530, 0.1339, 0.1256, 0.1583],\n",
       "            [0.2337, 0.1793, 0.1609, 0.1399, 0.1285, 0.1577],\n",
       "            [0.2352, 0.1738, 0.1668, 0.1574, 0.1477, 0.1191]],\n",
       "  \n",
       "           [[0.0964, 0.1466, 0.1572, 0.1823, 0.2121, 0.2054],\n",
       "            [0.0751, 0.1447, 0.1716, 0.2145, 0.2474, 0.1466],\n",
       "            [0.0752, 0.1525, 0.1764, 0.2163, 0.2495, 0.1301],\n",
       "            [0.0749, 0.1624, 0.1834, 0.2165, 0.2447, 0.1181],\n",
       "            [0.0728, 0.1712, 0.1910, 0.2170, 0.2367, 0.1114],\n",
       "            [0.0666, 0.1311, 0.1721, 0.2308, 0.2631, 0.1364]],\n",
       "  \n",
       "           [[0.1133, 0.1756, 0.1645, 0.1746, 0.1872, 0.1848],\n",
       "            [0.1409, 0.1734, 0.1630, 0.1638, 0.1645, 0.1944],\n",
       "            [0.1392, 0.1752, 0.1646, 0.1641, 0.1638, 0.1930],\n",
       "            [0.1364, 0.1754, 0.1657, 0.1628, 0.1603, 0.1994],\n",
       "            [0.1289, 0.1736, 0.1668, 0.1627, 0.1582, 0.2099],\n",
       "            [0.1363, 0.1719, 0.1778, 0.1864, 0.1916, 0.1359]],\n",
       "  \n",
       "           [[0.1348, 0.1582, 0.1739, 0.1956, 0.1907, 0.1469],\n",
       "            [0.2510, 0.1640, 0.1815, 0.1493, 0.1290, 0.1252],\n",
       "            [0.2487, 0.1689, 0.1846, 0.1482, 0.1236, 0.1259],\n",
       "            [0.2461, 0.1644, 0.1794, 0.1494, 0.1290, 0.1317],\n",
       "            [0.2522, 0.1453, 0.1568, 0.1499, 0.1518, 0.1440],\n",
       "            [0.1888, 0.1468, 0.1665, 0.1655, 0.1664, 0.1662]],\n",
       "  \n",
       "           [[0.1109, 0.1771, 0.1790, 0.1640, 0.1463, 0.2228],\n",
       "            [0.0953, 0.1550, 0.1576, 0.1484, 0.1315, 0.3122],\n",
       "            [0.0730, 0.1466, 0.1557, 0.1526, 0.1350, 0.3370],\n",
       "            [0.0719, 0.1495, 0.1580, 0.1526, 0.1322, 0.3358],\n",
       "            [0.0911, 0.1537, 0.1641, 0.1561, 0.1324, 0.3025],\n",
       "            [0.1671, 0.1565, 0.1692, 0.1631, 0.1375, 0.2065]],\n",
       "  \n",
       "           [[0.2260, 0.1331, 0.1535, 0.1807, 0.1882, 0.1186],\n",
       "            [0.1787, 0.1720, 0.1722, 0.1712, 0.1657, 0.1402],\n",
       "            [0.2000, 0.1754, 0.1684, 0.1620, 0.1577, 0.1365],\n",
       "            [0.2102, 0.1757, 0.1690, 0.1606, 0.1549, 0.1297],\n",
       "            [0.2091, 0.1683, 0.1687, 0.1631, 0.1553, 0.1355],\n",
       "            [0.1878, 0.1582, 0.1564, 0.1581, 0.1590, 0.1805]],\n",
       "  \n",
       "           [[0.1517, 0.1844, 0.1831, 0.1767, 0.1728, 0.1314],\n",
       "            [0.1646, 0.1714, 0.1519, 0.1517, 0.1621, 0.1984],\n",
       "            [0.1591, 0.1563, 0.1475, 0.1503, 0.1597, 0.2271],\n",
       "            [0.1542, 0.1499, 0.1469, 0.1524, 0.1647, 0.2318],\n",
       "            [0.1432, 0.1426, 0.1430, 0.1523, 0.1698, 0.2491],\n",
       "            [0.1526, 0.1558, 0.1530, 0.1506, 0.1543, 0.2337]],\n",
       "  \n",
       "           [[0.1953, 0.2014, 0.1683, 0.1590, 0.1569, 0.1192],\n",
       "            [0.1566, 0.2264, 0.2037, 0.1742, 0.1441, 0.0949],\n",
       "            [0.1391, 0.2303, 0.2108, 0.1784, 0.1459, 0.0955],\n",
       "            [0.1170, 0.2283, 0.2131, 0.1866, 0.1591, 0.0959],\n",
       "            [0.1188, 0.2188, 0.2041, 0.1882, 0.1708, 0.0993],\n",
       "            [0.1084, 0.1888, 0.1749, 0.1966, 0.2130, 0.1183]],\n",
       "  \n",
       "           [[0.2045, 0.1983, 0.1590, 0.1501, 0.1480, 0.1401],\n",
       "            [0.1317, 0.2428, 0.1582, 0.1267, 0.1224, 0.2183],\n",
       "            [0.1306, 0.2269, 0.1497, 0.1253, 0.1162, 0.2513],\n",
       "            [0.1303, 0.1298, 0.1180, 0.1180, 0.1015, 0.4024],\n",
       "            [0.1454, 0.0000, 0.0396, 0.0714, 0.0431, 0.7004],\n",
       "            [0.2635, 0.1584, 0.1307, 0.1435, 0.1461, 0.1578]],\n",
       "  \n",
       "           [[0.2770, 0.1070, 0.1195, 0.1525, 0.1562, 0.1877],\n",
       "            [0.3034, 0.1360, 0.1321, 0.1250, 0.1205, 0.1831],\n",
       "            [0.2974, 0.1372, 0.1343, 0.1220, 0.1113, 0.1977],\n",
       "            [0.2702, 0.1367, 0.1386, 0.1277, 0.1112, 0.2156],\n",
       "            [0.2330, 0.1322, 0.1401, 0.1389, 0.1214, 0.2344],\n",
       "            [0.3348, 0.0939, 0.0971, 0.1289, 0.1452, 0.2001]],\n",
       "  \n",
       "           [[0.3541, 0.1248, 0.1205, 0.1046, 0.0912, 0.2049],\n",
       "            [0.3880, 0.0842, 0.0760, 0.0640, 0.0511, 0.3368],\n",
       "            [0.3402, 0.1105, 0.1026, 0.0875, 0.0721, 0.2870],\n",
       "            [0.3037, 0.1357, 0.1262, 0.1068, 0.0901, 0.2376],\n",
       "            [0.2903, 0.1448, 0.1373, 0.1181, 0.1015, 0.2079],\n",
       "            [0.3661, 0.2022, 0.1219, 0.0627, 0.0830, 0.1641]],\n",
       "  \n",
       "           [[0.1604, 0.1934, 0.1788, 0.1929, 0.2147, 0.0599],\n",
       "            [0.1550, 0.2262, 0.1934, 0.1885, 0.1976, 0.0393],\n",
       "            [0.1346, 0.2087, 0.1949, 0.1988, 0.2131, 0.0499],\n",
       "            [0.1336, 0.2000, 0.1888, 0.1991, 0.2198, 0.0587],\n",
       "            [0.1319, 0.2043, 0.1885, 0.1971, 0.2198, 0.0584],\n",
       "            [0.2146, 0.1838, 0.1800, 0.1725, 0.1777, 0.0714]]]]),\n",
       "  tensor([[[[0.1756, 0.0555, 0.1123, 0.1594, 0.1726, 0.3247],\n",
       "            [0.1844, 0.1949, 0.1991, 0.1890, 0.1825, 0.0501],\n",
       "            [0.1480, 0.1907, 0.1992, 0.1983, 0.1966, 0.0671],\n",
       "            [0.1267, 0.1778, 0.1942, 0.2038, 0.2050, 0.0925],\n",
       "            [0.1292, 0.1626, 0.1795, 0.1949, 0.2088, 0.1251],\n",
       "            [0.1503, 0.0219, 0.0827, 0.1642, 0.2174, 0.3635]],\n",
       "  \n",
       "           [[0.1480, 0.2064, 0.1816, 0.1329, 0.0958, 0.2353],\n",
       "            [0.2009, 0.1484, 0.1446, 0.1461, 0.1448, 0.2153],\n",
       "            [0.1915, 0.1528, 0.1435, 0.1463, 0.1512, 0.2148],\n",
       "            [0.1865, 0.1535, 0.1391, 0.1385, 0.1453, 0.2371],\n",
       "            [0.1940, 0.1465, 0.1301, 0.1254, 0.1306, 0.2734],\n",
       "            [0.1795, 0.1476, 0.1439, 0.1456, 0.1449, 0.2386]],\n",
       "  \n",
       "           [[0.0423, 0.2081, 0.1965, 0.1800, 0.1600, 0.2131],\n",
       "            [0.1639, 0.1674, 0.1675, 0.1604, 0.1381, 0.2027],\n",
       "            [0.1466, 0.1751, 0.1675, 0.1607, 0.1437, 0.2065],\n",
       "            [0.1189, 0.1827, 0.1708, 0.1593, 0.1439, 0.2244],\n",
       "            [0.0962, 0.1813, 0.1813, 0.1676, 0.1441, 0.2296],\n",
       "            [0.1368, 0.1739, 0.1735, 0.1599, 0.1413, 0.2146]],\n",
       "  \n",
       "           [[0.1261, 0.1701, 0.1558, 0.1638, 0.1540, 0.2302],\n",
       "            [0.1235, 0.1941, 0.1758, 0.1585, 0.1353, 0.2128],\n",
       "            [0.1469, 0.2142, 0.1855, 0.1533, 0.1242, 0.1759],\n",
       "            [0.1565, 0.2120, 0.1903, 0.1627, 0.1343, 0.1443],\n",
       "            [0.1607, 0.1966, 0.1932, 0.1798, 0.1547, 0.1149],\n",
       "            [0.1757, 0.2009, 0.1779, 0.1617, 0.1464, 0.1374]],\n",
       "  \n",
       "           [[0.1661, 0.1845, 0.2021, 0.1958, 0.1552, 0.0963],\n",
       "            [0.1815, 0.1471, 0.1858, 0.1916, 0.1630, 0.1311],\n",
       "            [0.1509, 0.1528, 0.1941, 0.2055, 0.1810, 0.1157],\n",
       "            [0.1244, 0.1375, 0.1850, 0.2188, 0.2121, 0.1223],\n",
       "            [0.1101, 0.1219, 0.1736, 0.2279, 0.2375, 0.1291],\n",
       "            [0.1250, 0.1514, 0.1751, 0.1988, 0.1954, 0.1543]],\n",
       "  \n",
       "           [[0.1273, 0.1846, 0.1753, 0.1629, 0.1701, 0.1798],\n",
       "            [0.0833, 0.0899, 0.1084, 0.1214, 0.1232, 0.4739],\n",
       "            [0.0892, 0.0993, 0.1173, 0.1444, 0.1556, 0.3942],\n",
       "            [0.0932, 0.1068, 0.1196, 0.1491, 0.1683, 0.3631],\n",
       "            [0.0948, 0.1320, 0.1331, 0.1455, 0.1608, 0.3338],\n",
       "            [0.0675, 0.1147, 0.1522, 0.1514, 0.1355, 0.3787]],\n",
       "  \n",
       "           [[0.1372, 0.2409, 0.2212, 0.1846, 0.1429, 0.0733],\n",
       "            [0.1368, 0.1796, 0.1781, 0.1733, 0.1764, 0.1558],\n",
       "            [0.1385, 0.1772, 0.1769, 0.1739, 0.1783, 0.1552],\n",
       "            [0.1419, 0.1752, 0.1766, 0.1765, 0.1827, 0.1471],\n",
       "            [0.1462, 0.1748, 0.1773, 0.1797, 0.1871, 0.1349],\n",
       "            [0.1676, 0.1789, 0.1787, 0.1838, 0.1962, 0.0947]],\n",
       "  \n",
       "           [[0.1348, 0.2386, 0.2209, 0.1757, 0.1271, 0.1029],\n",
       "            [0.1187, 0.2076, 0.2054, 0.1778, 0.1488, 0.1416],\n",
       "            [0.1313, 0.1901, 0.1966, 0.1783, 0.1547, 0.1490],\n",
       "            [0.1367, 0.1752, 0.1854, 0.1759, 0.1611, 0.1658],\n",
       "            [0.1295, 0.1732, 0.1812, 0.1753, 0.1670, 0.1737],\n",
       "            [0.1490, 0.2060, 0.1763, 0.1385, 0.1164, 0.2138]],\n",
       "  \n",
       "           [[0.2921, 0.1876, 0.1597, 0.1451, 0.1288, 0.0867],\n",
       "            [0.2348, 0.2087, 0.1828, 0.1644, 0.1434, 0.0659],\n",
       "            [0.2526, 0.2116, 0.1766, 0.1588, 0.1400, 0.0604],\n",
       "            [0.2509, 0.2173, 0.1738, 0.1553, 0.1361, 0.0667],\n",
       "            [0.2372, 0.2220, 0.1753, 0.1541, 0.1317, 0.0797],\n",
       "            [0.1960, 0.2085, 0.1737, 0.1398, 0.1065, 0.1755]],\n",
       "  \n",
       "           [[0.0765, 0.2088, 0.2061, 0.1978, 0.1947, 0.1161],\n",
       "            [0.1461, 0.1545, 0.1749, 0.2073, 0.2382, 0.0791],\n",
       "            [0.1585, 0.1491, 0.1684, 0.2022, 0.2394, 0.0823],\n",
       "            [0.1508, 0.1508, 0.1716, 0.2062, 0.2418, 0.0788],\n",
       "            [0.1426, 0.1527, 0.1735, 0.2092, 0.2432, 0.0788],\n",
       "            [0.0702, 0.1622, 0.1847, 0.1833, 0.1658, 0.2337]],\n",
       "  \n",
       "           [[0.1480, 0.1664, 0.1802, 0.1977, 0.2040, 0.1037],\n",
       "            [0.2594, 0.1599, 0.1498, 0.1532, 0.1569, 0.1208],\n",
       "            [0.2584, 0.1577, 0.1515, 0.1539, 0.1543, 0.1241],\n",
       "            [0.2474, 0.1562, 0.1540, 0.1589, 0.1585, 0.1250],\n",
       "            [0.2244, 0.1576, 0.1590, 0.1687, 0.1702, 0.1202],\n",
       "            [0.2200, 0.1695, 0.1653, 0.1777, 0.1799, 0.0876]],\n",
       "  \n",
       "           [[0.0956, 0.1960, 0.1923, 0.1897, 0.1737, 0.1527],\n",
       "            [0.1185, 0.1550, 0.1767, 0.1959, 0.1921, 0.1617],\n",
       "            [0.1163, 0.1401, 0.1665, 0.2000, 0.2034, 0.1736],\n",
       "            [0.1277, 0.1411, 0.1588, 0.1941, 0.2017, 0.1765],\n",
       "            [0.1359, 0.1647, 0.1632, 0.1850, 0.1879, 0.1634],\n",
       "            [0.1628, 0.1689, 0.1591, 0.1593, 0.1571, 0.1928]]]]),\n",
       "  tensor([[[[0.2856, 0.0619, 0.0493, 0.1048, 0.1152, 0.3832],\n",
       "            [0.2068, 0.1190, 0.1165, 0.1674, 0.1264, 0.2639],\n",
       "            [0.2159, 0.1174, 0.1053, 0.1534, 0.1225, 0.2855],\n",
       "            [0.2493, 0.1224, 0.0989, 0.1295, 0.1053, 0.2947],\n",
       "            [0.2743, 0.1358, 0.1075, 0.1143, 0.0908, 0.2772],\n",
       "            [0.1998, 0.1367, 0.1083, 0.1151, 0.1138, 0.3263]],\n",
       "  \n",
       "           [[0.2844, 0.1663, 0.1345, 0.1307, 0.1513, 0.1328],\n",
       "            [0.2376, 0.1723, 0.1501, 0.1482, 0.1606, 0.1313],\n",
       "            [0.2921, 0.1598, 0.1391, 0.1352, 0.1459, 0.1279],\n",
       "            [0.3484, 0.1479, 0.1280, 0.1227, 0.1291, 0.1240],\n",
       "            [0.3707, 0.1487, 0.1247, 0.1199, 0.1270, 0.1090],\n",
       "            [0.2247, 0.1719, 0.1552, 0.1571, 0.1647, 0.1264]],\n",
       "  \n",
       "           [[0.1195, 0.1556, 0.1609, 0.1767, 0.2040, 0.1832],\n",
       "            [0.1684, 0.1222, 0.1431, 0.1707, 0.2000, 0.1955],\n",
       "            [0.1776, 0.1236, 0.1430, 0.1707, 0.2009, 0.1843],\n",
       "            [0.1828, 0.1269, 0.1434, 0.1682, 0.1932, 0.1855],\n",
       "            [0.1850, 0.1299, 0.1419, 0.1623, 0.1811, 0.1999],\n",
       "            [0.1262, 0.1463, 0.1428, 0.1568, 0.1765, 0.2514]],\n",
       "  \n",
       "           [[0.2408, 0.1523, 0.1384, 0.1366, 0.1479, 0.1841],\n",
       "            [0.2166, 0.1338, 0.1412, 0.1608, 0.1902, 0.1574],\n",
       "            [0.2008, 0.1303, 0.1432, 0.1697, 0.1983, 0.1576],\n",
       "            [0.1887, 0.1258, 0.1433, 0.1767, 0.2067, 0.1587],\n",
       "            [0.1857, 0.1227, 0.1416, 0.1788, 0.2115, 0.1597],\n",
       "            [0.1416, 0.1362, 0.1517, 0.1997, 0.2344, 0.1363]],\n",
       "  \n",
       "           [[0.2214, 0.0761, 0.0985, 0.1317, 0.1709, 0.3014],\n",
       "            [0.1677, 0.0449, 0.0875, 0.1270, 0.1589, 0.4140],\n",
       "            [0.1667, 0.0529, 0.0928, 0.1293, 0.1594, 0.3989],\n",
       "            [0.1776, 0.0575, 0.0922, 0.1259, 0.1568, 0.3899],\n",
       "            [0.1842, 0.0472, 0.0851, 0.1226, 0.1575, 0.4033],\n",
       "            [0.2629, 0.1004, 0.0914, 0.0997, 0.1267, 0.3188]],\n",
       "  \n",
       "           [[0.1724, 0.1820, 0.1966, 0.1940, 0.1719, 0.0831],\n",
       "            [0.1264, 0.1890, 0.2026, 0.2002, 0.1806, 0.1014],\n",
       "            [0.1269, 0.1901, 0.1983, 0.1935, 0.1768, 0.1144],\n",
       "            [0.1249, 0.1938, 0.1989, 0.1915, 0.1757, 0.1153],\n",
       "            [0.1185, 0.1957, 0.2054, 0.2001, 0.1842, 0.0962],\n",
       "            [0.0830, 0.1869, 0.2044, 0.2102, 0.2054, 0.1101]],\n",
       "  \n",
       "           [[0.1821, 0.2001, 0.1587, 0.1253, 0.1318, 0.2020],\n",
       "            [0.1084, 0.1195, 0.1106, 0.1331, 0.1823, 0.3461],\n",
       "            [0.1064, 0.1339, 0.1131, 0.1267, 0.1727, 0.3472],\n",
       "            [0.1280, 0.1544, 0.1186, 0.1169, 0.1540, 0.3280],\n",
       "            [0.1466, 0.1727, 0.1264, 0.1107, 0.1369, 0.3067],\n",
       "            [0.1184, 0.2096, 0.1825, 0.1485, 0.1426, 0.1984]],\n",
       "  \n",
       "           [[0.1871, 0.1511, 0.1554, 0.1466, 0.1237, 0.2362],\n",
       "            [0.1783, 0.1367, 0.1437, 0.1499, 0.1452, 0.2462],\n",
       "            [0.1799, 0.1402, 0.1459, 0.1495, 0.1436, 0.2409],\n",
       "            [0.1843, 0.1462, 0.1498, 0.1473, 0.1360, 0.2365],\n",
       "            [0.1867, 0.1494, 0.1522, 0.1454, 0.1298, 0.2365],\n",
       "            [0.2135, 0.1485, 0.1429, 0.1271, 0.1011, 0.2669]],\n",
       "  \n",
       "           [[0.2523, 0.1107, 0.1095, 0.1535, 0.1930, 0.1810],\n",
       "            [0.1329, 0.1645, 0.1830, 0.1995, 0.1983, 0.1218],\n",
       "            [0.1298, 0.1182, 0.1539, 0.2100, 0.2398, 0.1483],\n",
       "            [0.1347, 0.0970, 0.1361, 0.2076, 0.2565, 0.1681],\n",
       "            [0.1679, 0.0905, 0.1243, 0.1942, 0.2478, 0.1753],\n",
       "            [0.2400, 0.0946, 0.0844, 0.1295, 0.1736, 0.2778]],\n",
       "  \n",
       "           [[0.1986, 0.1334, 0.1293, 0.1412, 0.1630, 0.2345],\n",
       "            [0.2321, 0.1441, 0.1438, 0.1416, 0.1389, 0.1995],\n",
       "            [0.2548, 0.1354, 0.1374, 0.1386, 0.1381, 0.1957],\n",
       "            [0.2675, 0.1230, 0.1251, 0.1356, 0.1478, 0.2010],\n",
       "            [0.2604, 0.1124, 0.1135, 0.1348, 0.1637, 0.2151],\n",
       "            [0.2853, 0.0429, 0.0572, 0.1003, 0.1452, 0.3690]],\n",
       "  \n",
       "           [[0.1458, 0.1730, 0.1904, 0.2046, 0.2063, 0.0799],\n",
       "            [0.4861, 0.0558, 0.1324, 0.1512, 0.0795, 0.0950],\n",
       "            [0.4500, 0.0647, 0.1291, 0.1517, 0.1032, 0.1013],\n",
       "            [0.4477, 0.0607, 0.1252, 0.1559, 0.1227, 0.0877],\n",
       "            [0.4220, 0.0614, 0.1264, 0.1592, 0.1305, 0.1005],\n",
       "            [0.0813, 0.1676, 0.1950, 0.2113, 0.2145, 0.1303]],\n",
       "  \n",
       "           [[0.1960, 0.1602, 0.1541, 0.1541, 0.1470, 0.1887],\n",
       "            [0.1353, 0.2141, 0.2019, 0.1837, 0.1591, 0.1059],\n",
       "            [0.1259, 0.2179, 0.2080, 0.1881, 0.1590, 0.1011],\n",
       "            [0.1245, 0.2145, 0.2057, 0.1896, 0.1632, 0.1025],\n",
       "            [0.1254, 0.2085, 0.1993, 0.1892, 0.1701, 0.1073],\n",
       "            [0.1345, 0.1822, 0.1769, 0.1761, 0.1554, 0.1749]]]]),\n",
       "  tensor([[[[0.1677, 0.2107, 0.2095, 0.1898, 0.1594, 0.0630],\n",
       "            [0.1382, 0.2296, 0.2322, 0.2049, 0.1492, 0.0460],\n",
       "            [0.1234, 0.2106, 0.2260, 0.2128, 0.1676, 0.0595],\n",
       "            [0.1142, 0.1958, 0.2166, 0.2156, 0.1901, 0.0676],\n",
       "            [0.1192, 0.1924, 0.2105, 0.2123, 0.1967, 0.0688],\n",
       "            [0.1532, 0.1921, 0.1977, 0.1930, 0.1835, 0.0805]],\n",
       "  \n",
       "           [[0.2447, 0.1989, 0.1934, 0.1604, 0.1482, 0.0544],\n",
       "            [0.2782, 0.1627, 0.1659, 0.1391, 0.1308, 0.1234],\n",
       "            [0.3292, 0.1234, 0.1365, 0.1314, 0.1684, 0.1110],\n",
       "            [0.5969, 0.0134, 0.0457, 0.0925, 0.2515, 0.0000],\n",
       "            [0.3897, 0.0568, 0.0770, 0.1039, 0.2498, 0.1229],\n",
       "            [0.1471, 0.1733, 0.1585, 0.1480, 0.1652, 0.2080]],\n",
       "  \n",
       "           [[0.2069, 0.1783, 0.1954, 0.1715, 0.1436, 0.1043],\n",
       "            [0.1845, 0.1202, 0.1628, 0.1913, 0.1734, 0.1678],\n",
       "            [0.1656, 0.1066, 0.1501, 0.1936, 0.1928, 0.1912],\n",
       "            [0.1450, 0.1031, 0.1351, 0.1847, 0.2118, 0.2203],\n",
       "            [0.1517, 0.1204, 0.1401, 0.1752, 0.2080, 0.2047],\n",
       "            [0.1882, 0.1373, 0.1494, 0.1697, 0.1824, 0.1730]],\n",
       "  \n",
       "           [[0.1514, 0.2520, 0.2174, 0.1828, 0.1553, 0.0410],\n",
       "            [0.2105, 0.1710, 0.1383, 0.1543, 0.1525, 0.1734],\n",
       "            [0.2295, 0.1815, 0.1315, 0.1421, 0.1630, 0.1524],\n",
       "            [0.2264, 0.1930, 0.1392, 0.1462, 0.1724, 0.1227],\n",
       "            [0.2103, 0.1955, 0.1533, 0.1632, 0.1742, 0.1035],\n",
       "            [0.1917, 0.2072, 0.1477, 0.1471, 0.1869, 0.1195]],\n",
       "  \n",
       "           [[0.3013, 0.1605, 0.1427, 0.1416, 0.1452, 0.1088],\n",
       "            [0.1787, 0.1829, 0.1762, 0.1760, 0.1741, 0.1121],\n",
       "            [0.1710, 0.1860, 0.1779, 0.1774, 0.1801, 0.1077],\n",
       "            [0.1443, 0.2034, 0.1887, 0.1856, 0.1953, 0.0827],\n",
       "            [0.1158, 0.2190, 0.1997, 0.1943, 0.2096, 0.0617],\n",
       "            [0.3120, 0.1277, 0.1139, 0.1178, 0.1485, 0.1802]],\n",
       "  \n",
       "           [[0.2110, 0.2301, 0.1947, 0.1514, 0.1161, 0.0967],\n",
       "            [0.0994, 0.2511, 0.2097, 0.1914, 0.1467, 0.1016],\n",
       "            [0.1246, 0.2249, 0.2036, 0.1858, 0.1494, 0.1117],\n",
       "            [0.1461, 0.2108, 0.1961, 0.1783, 0.1478, 0.1210],\n",
       "            [0.1627, 0.2032, 0.1881, 0.1708, 0.1446, 0.1306],\n",
       "            [0.1010, 0.2304, 0.2162, 0.1819, 0.1565, 0.1140]],\n",
       "  \n",
       "           [[0.1426, 0.1662, 0.1749, 0.1791, 0.1766, 0.1607],\n",
       "            [0.2598, 0.1317, 0.1464, 0.1480, 0.1288, 0.1852],\n",
       "            [0.2580, 0.1296, 0.1461, 0.1517, 0.1362, 0.1784],\n",
       "            [0.2500, 0.1303, 0.1487, 0.1571, 0.1435, 0.1705],\n",
       "            [0.2397, 0.1327, 0.1528, 0.1636, 0.1512, 0.1601],\n",
       "            [0.1660, 0.1620, 0.1599, 0.1579, 0.1598, 0.1943]],\n",
       "  \n",
       "           [[0.1238, 0.1479, 0.1659, 0.1797, 0.1945, 0.1883],\n",
       "            [0.1549, 0.1236, 0.1548, 0.1802, 0.1994, 0.1871],\n",
       "            [0.1617, 0.1240, 0.1482, 0.1736, 0.1966, 0.1958],\n",
       "            [0.1763, 0.1322, 0.1429, 0.1621, 0.1854, 0.2012],\n",
       "            [0.1897, 0.1449, 0.1427, 0.1528, 0.1720, 0.1979],\n",
       "            [0.1827, 0.1415, 0.1552, 0.1642, 0.1767, 0.1797]],\n",
       "  \n",
       "           [[0.0982, 0.1149, 0.1150, 0.1328, 0.1381, 0.4010],\n",
       "            [0.0599, 0.1522, 0.1420, 0.1317, 0.1253, 0.3888],\n",
       "            [0.0762, 0.1573, 0.1424, 0.1306, 0.1263, 0.3672],\n",
       "            [0.1082, 0.1672, 0.1433, 0.1298, 0.1188, 0.3328],\n",
       "            [0.1266, 0.1874, 0.1638, 0.1433, 0.1127, 0.2661],\n",
       "            [0.2093, 0.1789, 0.1611, 0.1370, 0.0961, 0.2177]],\n",
       "  \n",
       "           [[0.1143, 0.1421, 0.1940, 0.2004, 0.1679, 0.1813],\n",
       "            [0.1448, 0.1623, 0.1628, 0.1705, 0.1652, 0.1943],\n",
       "            [0.1511, 0.1672, 0.1583, 0.1629, 0.1619, 0.1986],\n",
       "            [0.1563, 0.1744, 0.1607, 0.1560, 0.1538, 0.1988],\n",
       "            [0.1591, 0.1799, 0.1650, 0.1545, 0.1504, 0.1910],\n",
       "            [0.1610, 0.1429, 0.1626, 0.1753, 0.1795, 0.1787]],\n",
       "  \n",
       "           [[0.2274, 0.1182, 0.1450, 0.1627, 0.1673, 0.1794],\n",
       "            [0.2082, 0.1601, 0.1710, 0.1718, 0.1639, 0.1250],\n",
       "            [0.2228, 0.1689, 0.1708, 0.1590, 0.1427, 0.1358],\n",
       "            [0.2236, 0.1675, 0.1699, 0.1574, 0.1403, 0.1413],\n",
       "            [0.2108, 0.1573, 0.1681, 0.1664, 0.1563, 0.1411],\n",
       "            [0.2568, 0.1460, 0.1614, 0.1661, 0.1557, 0.1140]],\n",
       "  \n",
       "           [[0.2707, 0.1632, 0.1524, 0.1449, 0.1435, 0.1252],\n",
       "            [0.1730, 0.1715, 0.1843, 0.1759, 0.1497, 0.1456],\n",
       "            [0.1364, 0.1546, 0.1801, 0.1957, 0.1849, 0.1483],\n",
       "            [0.1149, 0.1385, 0.1723, 0.2078, 0.2108, 0.1557],\n",
       "            [0.1323, 0.1325, 0.1596, 0.1975, 0.2096, 0.1684],\n",
       "            [0.1439, 0.1826, 0.1700, 0.1608, 0.1390, 0.2037]]]]),\n",
       "  tensor([[[[0.2118, 0.1887, 0.1697, 0.1590, 0.1479, 0.1228],\n",
       "            [0.2154, 0.1795, 0.1659, 0.1594, 0.1385, 0.1413],\n",
       "            [0.2061, 0.1864, 0.1627, 0.1577, 0.1406, 0.1465],\n",
       "            [0.2022, 0.1926, 0.1609, 0.1545, 0.1372, 0.1526],\n",
       "            [0.1938, 0.1888, 0.1633, 0.1577, 0.1366, 0.1598],\n",
       "            [0.1364, 0.1836, 0.1821, 0.1836, 0.1730, 0.1413]],\n",
       "  \n",
       "           [[0.1911, 0.2078, 0.1520, 0.1484, 0.1637, 0.1370],\n",
       "            [0.2233, 0.2209, 0.1581, 0.1411, 0.1458, 0.1108],\n",
       "            [0.2240, 0.2146, 0.1659, 0.1536, 0.1468, 0.0950],\n",
       "            [0.1918, 0.1903, 0.1699, 0.1794, 0.1807, 0.0879],\n",
       "            [0.1464, 0.1718, 0.1671, 0.1961, 0.2191, 0.0995],\n",
       "            [0.1526, 0.1636, 0.1475, 0.1725, 0.1960, 0.1678]],\n",
       "  \n",
       "           [[0.0887, 0.1482, 0.1478, 0.1568, 0.1482, 0.3103],\n",
       "            [0.2138, 0.1277, 0.0985, 0.1058, 0.0991, 0.3551],\n",
       "            [0.2528, 0.1236, 0.1072, 0.1001, 0.0892, 0.3271],\n",
       "            [0.2760, 0.1181, 0.1199, 0.1091, 0.0924, 0.2846],\n",
       "            [0.2679, 0.1178, 0.1351, 0.1282, 0.1059, 0.2450],\n",
       "            [0.2109, 0.1253, 0.1595, 0.1421, 0.1152, 0.2470]],\n",
       "  \n",
       "           [[0.1341, 0.1338, 0.1748, 0.2119, 0.2301, 0.1153],\n",
       "            [0.1641, 0.1365, 0.1596, 0.1827, 0.1985, 0.1586],\n",
       "            [0.1815, 0.1322, 0.1535, 0.1801, 0.1969, 0.1558],\n",
       "            [0.1939, 0.1352, 0.1515, 0.1762, 0.1914, 0.1518],\n",
       "            [0.1813, 0.1415, 0.1565, 0.1779, 0.1918, 0.1511],\n",
       "            [0.1826, 0.1407, 0.1634, 0.1861, 0.1934, 0.1339]],\n",
       "  \n",
       "           [[0.2646, 0.1362, 0.1198, 0.1172, 0.1368, 0.2254],\n",
       "            [0.1875, 0.1568, 0.1531, 0.1496, 0.1548, 0.1982],\n",
       "            [0.1754, 0.1639, 0.1616, 0.1574, 0.1591, 0.1827],\n",
       "            [0.1762, 0.1687, 0.1662, 0.1611, 0.1615, 0.1663],\n",
       "            [0.1875, 0.1707, 0.1656, 0.1596, 0.1615, 0.1550],\n",
       "            [0.2215, 0.1600, 0.1488, 0.1420, 0.1487, 0.1791]],\n",
       "  \n",
       "           [[0.2484, 0.1546, 0.1345, 0.1312, 0.1523, 0.1789],\n",
       "            [0.1707, 0.1811, 0.1670, 0.1561, 0.1517, 0.1733],\n",
       "            [0.1299, 0.1854, 0.1839, 0.1749, 0.1617, 0.1641],\n",
       "            [0.1092, 0.1800, 0.1882, 0.1856, 0.1706, 0.1664],\n",
       "            [0.1100, 0.1804, 0.1857, 0.1838, 0.1721, 0.1681],\n",
       "            [0.1604, 0.1736, 0.1767, 0.1767, 0.1789, 0.1337]],\n",
       "  \n",
       "           [[0.1490, 0.1249, 0.1239, 0.1752, 0.2205, 0.2065],\n",
       "            [0.1455, 0.1541, 0.1598, 0.1812, 0.1929, 0.1665],\n",
       "            [0.1388, 0.1503, 0.1503, 0.1817, 0.2068, 0.1720],\n",
       "            [0.1399, 0.1506, 0.1449, 0.1775, 0.2095, 0.1777],\n",
       "            [0.1431, 0.1541, 0.1466, 0.1750, 0.2059, 0.1753],\n",
       "            [0.1625, 0.1662, 0.1543, 0.1716, 0.1788, 0.1665]],\n",
       "  \n",
       "           [[0.1515, 0.1845, 0.1706, 0.1642, 0.1652, 0.1641],\n",
       "            [0.1200, 0.1900, 0.1991, 0.2031, 0.1889, 0.0989],\n",
       "            [0.1155, 0.1929, 0.1846, 0.1854, 0.1912, 0.1304],\n",
       "            [0.1083, 0.1863, 0.1920, 0.1914, 0.1923, 0.1297],\n",
       "            [0.1174, 0.2094, 0.2036, 0.1794, 0.1653, 0.1250],\n",
       "            [0.0865, 0.1832, 0.1906, 0.1885, 0.1850, 0.1663]],\n",
       "  \n",
       "           [[0.0971, 0.1474, 0.1884, 0.1949, 0.1722, 0.1998],\n",
       "            [0.1240, 0.1011, 0.1845, 0.2283, 0.2178, 0.1442],\n",
       "            [0.1514, 0.1269, 0.1896, 0.2156, 0.1800, 0.1366],\n",
       "            [0.1698, 0.1443, 0.1950, 0.2061, 0.1582, 0.1267],\n",
       "            [0.1917, 0.1544, 0.1962, 0.1981, 0.1446, 0.1149],\n",
       "            [0.1583, 0.1780, 0.1918, 0.1885, 0.1529, 0.1306]],\n",
       "  \n",
       "           [[0.1469, 0.1724, 0.1734, 0.1691, 0.1618, 0.1763],\n",
       "            [0.1479, 0.1458, 0.1511, 0.1631, 0.1699, 0.2221],\n",
       "            [0.1620, 0.1391, 0.1260, 0.1433, 0.1756, 0.2540],\n",
       "            [0.1684, 0.1596, 0.1312, 0.1345, 0.1649, 0.2414],\n",
       "            [0.1572, 0.1961, 0.1655, 0.1432, 0.1415, 0.1965],\n",
       "            [0.1593, 0.1811, 0.1734, 0.1568, 0.1440, 0.1855]],\n",
       "  \n",
       "           [[0.0981, 0.1691, 0.1910, 0.2114, 0.2222, 0.1082],\n",
       "            [0.1287, 0.1913, 0.1704, 0.1887, 0.2135, 0.1074],\n",
       "            [0.1482, 0.1627, 0.1645, 0.1869, 0.2102, 0.1275],\n",
       "            [0.1689, 0.1276, 0.1554, 0.1832, 0.2101, 0.1548],\n",
       "            [0.1747, 0.1192, 0.1475, 0.1782, 0.2127, 0.1678],\n",
       "            [0.1711, 0.1239, 0.1600, 0.1811, 0.2011, 0.1628]],\n",
       "  \n",
       "           [[0.2483, 0.2218, 0.1622, 0.1211, 0.1432, 0.1034],\n",
       "            [0.2380, 0.1963, 0.1796, 0.1483, 0.1385, 0.0993],\n",
       "            [0.2526, 0.1976, 0.1842, 0.1414, 0.1292, 0.0950],\n",
       "            [0.2423, 0.1897, 0.1886, 0.1512, 0.1403, 0.0879],\n",
       "            [0.2209, 0.1765, 0.1868, 0.1664, 0.1584, 0.0910],\n",
       "            [0.1809, 0.1748, 0.1812, 0.1818, 0.1738, 0.1075]]]]),\n",
       "  tensor([[[[0.2157, 0.1815, 0.1762, 0.1598, 0.1300, 0.1369],\n",
       "            [0.2004, 0.1698, 0.1704, 0.1703, 0.1559, 0.1332],\n",
       "            [0.2157, 0.1727, 0.1613, 0.1577, 0.1488, 0.1438],\n",
       "            [0.2279, 0.1829, 0.1616, 0.1457, 0.1346, 0.1474],\n",
       "            [0.2288, 0.1922, 0.1708, 0.1433, 0.1232, 0.1417],\n",
       "            [0.1901, 0.1946, 0.1940, 0.1699, 0.1392, 0.1123]],\n",
       "  \n",
       "           [[0.1615, 0.1600, 0.1601, 0.1658, 0.1695, 0.1831],\n",
       "            [0.0683, 0.1961, 0.2061, 0.2059, 0.1914, 0.1321],\n",
       "            [0.0847, 0.1869, 0.1987, 0.2011, 0.1918, 0.1367],\n",
       "            [0.0983, 0.1805, 0.1920, 0.1989, 0.1978, 0.1326],\n",
       "            [0.1320, 0.1726, 0.1768, 0.1874, 0.2002, 0.1310],\n",
       "            [0.1902, 0.1543, 0.1533, 0.1695, 0.1942, 0.1385]],\n",
       "  \n",
       "           [[0.1564, 0.2047, 0.1750, 0.1354, 0.1038, 0.2247],\n",
       "            [0.0906, 0.1731, 0.1973, 0.1782, 0.1382, 0.2227],\n",
       "            [0.0885, 0.1620, 0.1932, 0.1763, 0.1283, 0.2516],\n",
       "            [0.1013, 0.1531, 0.1882, 0.1754, 0.1237, 0.2584],\n",
       "            [0.1224, 0.1522, 0.1882, 0.1802, 0.1273, 0.2297],\n",
       "            [0.1722, 0.2070, 0.1798, 0.1348, 0.0818, 0.2244]],\n",
       "  \n",
       "           [[0.1439, 0.1122, 0.1358, 0.1797, 0.2233, 0.2051],\n",
       "            [0.1758, 0.2011, 0.1778, 0.1576, 0.1325, 0.1552],\n",
       "            [0.1806, 0.1817, 0.1643, 0.1587, 0.1354, 0.1794],\n",
       "            [0.2067, 0.1491, 0.1208, 0.1297, 0.1376, 0.2561],\n",
       "            [0.2122, 0.1449, 0.1090, 0.1159, 0.1392, 0.2787],\n",
       "            [0.1723, 0.1717, 0.1154, 0.1145, 0.1544, 0.2719]],\n",
       "  \n",
       "           [[0.0130, 0.1434, 0.1725, 0.2161, 0.2089, 0.2460],\n",
       "            [0.0507, 0.0918, 0.1450, 0.1956, 0.2018, 0.3151],\n",
       "            [0.0644, 0.0842, 0.1533, 0.2065, 0.2011, 0.2906],\n",
       "            [0.1178, 0.0523, 0.1172, 0.2073, 0.2182, 0.2871],\n",
       "            [0.1651, 0.0878, 0.1156, 0.1905, 0.2056, 0.2354],\n",
       "            [0.1523, 0.1562, 0.1487, 0.1654, 0.1762, 0.2011]],\n",
       "  \n",
       "           [[0.1777, 0.1617, 0.1490, 0.1444, 0.1558, 0.2114],\n",
       "            [0.1502, 0.1646, 0.1602, 0.1676, 0.1748, 0.1827],\n",
       "            [0.1232, 0.1694, 0.1700, 0.1806, 0.1865, 0.1703],\n",
       "            [0.1190, 0.1697, 0.1696, 0.1801, 0.1884, 0.1731],\n",
       "            [0.1357, 0.1695, 0.1640, 0.1700, 0.1787, 0.1820],\n",
       "            [0.1617, 0.1714, 0.1484, 0.1447, 0.1633, 0.2105]],\n",
       "  \n",
       "           [[0.1050, 0.1179, 0.1564, 0.1665, 0.1394, 0.3147],\n",
       "            [0.1204, 0.1806, 0.1989, 0.1857, 0.1394, 0.1751],\n",
       "            [0.1453, 0.1988, 0.2026, 0.1815, 0.1346, 0.1372],\n",
       "            [0.1651, 0.2045, 0.2009, 0.1750, 0.1293, 0.1252],\n",
       "            [0.1656, 0.2201, 0.2100, 0.1723, 0.1168, 0.1152],\n",
       "            [0.1545, 0.2172, 0.2050, 0.1657, 0.1138, 0.1438]],\n",
       "  \n",
       "           [[0.2632, 0.1397, 0.0687, 0.1000, 0.1346, 0.2937],\n",
       "            [0.2608, 0.1516, 0.0922, 0.0887, 0.1183, 0.2884],\n",
       "            [0.2318, 0.1744, 0.1168, 0.0931, 0.1120, 0.2719],\n",
       "            [0.2030, 0.1800, 0.1375, 0.1065, 0.1179, 0.2551],\n",
       "            [0.1919, 0.1598, 0.1424, 0.1248, 0.1394, 0.2418],\n",
       "            [0.1590, 0.1602, 0.1559, 0.1255, 0.1700, 0.2294]],\n",
       "  \n",
       "           [[0.1170, 0.1950, 0.1986, 0.1848, 0.1820, 0.1226],\n",
       "            [0.1212, 0.1957, 0.2042, 0.1845, 0.1784, 0.1160],\n",
       "            [0.1046, 0.1905, 0.1974, 0.1864, 0.1842, 0.1369],\n",
       "            [0.0838, 0.1880, 0.1904, 0.1902, 0.1928, 0.1547],\n",
       "            [0.0882, 0.1925, 0.1888, 0.1876, 0.1896, 0.1532],\n",
       "            [0.1305, 0.2076, 0.1924, 0.1746, 0.1676, 0.1274]],\n",
       "  \n",
       "           [[0.1957, 0.1864, 0.1727, 0.1586, 0.1547, 0.1318],\n",
       "            [0.1346, 0.2531, 0.2197, 0.1796, 0.1655, 0.0476],\n",
       "            [0.1551, 0.2312, 0.1950, 0.1670, 0.1625, 0.0892],\n",
       "            [0.1819, 0.1907, 0.1573, 0.1559, 0.1727, 0.1415],\n",
       "            [0.2204, 0.1606, 0.1305, 0.1420, 0.1716, 0.1748],\n",
       "            [0.2245, 0.1670, 0.1423, 0.1357, 0.1480, 0.1823]],\n",
       "  \n",
       "           [[0.2060, 0.1362, 0.1506, 0.1685, 0.1772, 0.1616],\n",
       "            [0.2044, 0.1385, 0.1646, 0.1715, 0.1580, 0.1630],\n",
       "            [0.2021, 0.1476, 0.1727, 0.1755, 0.1586, 0.1435],\n",
       "            [0.2006, 0.1551, 0.1785, 0.1785, 0.1594, 0.1278],\n",
       "            [0.2051, 0.1558, 0.1764, 0.1746, 0.1537, 0.1346],\n",
       "            [0.1648, 0.1723, 0.1763, 0.1717, 0.1598, 0.1550]],\n",
       "  \n",
       "           [[0.1790, 0.1494, 0.1509, 0.1466, 0.1455, 0.2287],\n",
       "            [0.1704, 0.1385, 0.1533, 0.1563, 0.1632, 0.2183],\n",
       "            [0.1714, 0.1405, 0.1525, 0.1546, 0.1625, 0.2184],\n",
       "            [0.1702, 0.1452, 0.1536, 0.1553, 0.1647, 0.2109],\n",
       "            [0.1652, 0.1449, 0.1530, 0.1578, 0.1704, 0.2087],\n",
       "            [0.1752, 0.1484, 0.1544, 0.1566, 0.1603, 0.2052]]]]),\n",
       "  tensor([[[[0.1076, 0.1741, 0.1630, 0.1821, 0.1976, 0.1755],\n",
       "            [0.1085, 0.1842, 0.1724, 0.1851, 0.1964, 0.1533],\n",
       "            [0.1018, 0.1715, 0.1701, 0.1913, 0.2027, 0.1626],\n",
       "            [0.0973, 0.1585, 0.1665, 0.1969, 0.2102, 0.1706],\n",
       "            [0.0965, 0.1512, 0.1646, 0.2001, 0.2156, 0.1720],\n",
       "            [0.0924, 0.1652, 0.1549, 0.1900, 0.2104, 0.1871]],\n",
       "  \n",
       "           [[0.2364, 0.1517, 0.1478, 0.1383, 0.1404, 0.1855],\n",
       "            [0.2324, 0.1532, 0.1392, 0.1344, 0.1519, 0.1889],\n",
       "            [0.2147, 0.1536, 0.1415, 0.1432, 0.1695, 0.1775],\n",
       "            [0.2166, 0.1607, 0.1444, 0.1433, 0.1716, 0.1633],\n",
       "            [0.2314, 0.1693, 0.1468, 0.1370, 0.1614, 0.1542],\n",
       "            [0.2366, 0.1534, 0.1450, 0.1411, 0.1596, 0.1643]],\n",
       "  \n",
       "           [[0.1541, 0.1195, 0.1041, 0.1200, 0.1891, 0.3133],\n",
       "            [0.1014, 0.1289, 0.1305, 0.1467, 0.1980, 0.2945],\n",
       "            [0.1098, 0.1496, 0.1373, 0.1387, 0.1896, 0.2749],\n",
       "            [0.1075, 0.1581, 0.1419, 0.1374, 0.1842, 0.2710],\n",
       "            [0.1006, 0.1361, 0.1340, 0.1449, 0.1859, 0.2986],\n",
       "            [0.0519, 0.1315, 0.1574, 0.1728, 0.1960, 0.2904]],\n",
       "  \n",
       "           [[0.1803, 0.2130, 0.1926, 0.1645, 0.1502, 0.0993],\n",
       "            [0.1916, 0.2198, 0.1914, 0.1553, 0.1382, 0.1036],\n",
       "            [0.1851, 0.2168, 0.1905, 0.1567, 0.1402, 0.1106],\n",
       "            [0.1797, 0.2100, 0.1862, 0.1577, 0.1441, 0.1223],\n",
       "            [0.1770, 0.2051, 0.1817, 0.1572, 0.1465, 0.1326],\n",
       "            [0.1550, 0.1959, 0.1776, 0.1612, 0.1556, 0.1548]],\n",
       "  \n",
       "           [[0.1808, 0.1447, 0.1394, 0.1463, 0.1608, 0.2281],\n",
       "            [0.1899, 0.1624, 0.1491, 0.1454, 0.1499, 0.2034],\n",
       "            [0.1938, 0.1616, 0.1490, 0.1454, 0.1495, 0.2007],\n",
       "            [0.1823, 0.1582, 0.1493, 0.1481, 0.1529, 0.2092],\n",
       "            [0.1721, 0.1548, 0.1481, 0.1463, 0.1520, 0.2268],\n",
       "            [0.1168, 0.1713, 0.1639, 0.1574, 0.1556, 0.2350]],\n",
       "  \n",
       "           [[0.2726, 0.2149, 0.1856, 0.1288, 0.1064, 0.0917],\n",
       "            [0.2236, 0.1918, 0.1761, 0.1431, 0.1348, 0.1305],\n",
       "            [0.1329, 0.1293, 0.1442, 0.1593, 0.2086, 0.2256],\n",
       "            [0.1188, 0.0938, 0.1116, 0.1483, 0.2274, 0.3003],\n",
       "            [0.1321, 0.0804, 0.0910, 0.1305, 0.2213, 0.3447],\n",
       "            [0.2737, 0.0334, 0.0268, 0.0634, 0.1641, 0.4387]],\n",
       "  \n",
       "           [[0.1244, 0.1362, 0.1631, 0.1942, 0.2165, 0.1656],\n",
       "            [0.1459, 0.1235, 0.1446, 0.1887, 0.2301, 0.1672],\n",
       "            [0.1600, 0.1220, 0.1363, 0.1807, 0.2287, 0.1723],\n",
       "            [0.1644, 0.1299, 0.1381, 0.1787, 0.2265, 0.1623],\n",
       "            [0.1469, 0.1359, 0.1497, 0.1903, 0.2334, 0.1438],\n",
       "            [0.1353, 0.1482, 0.1659, 0.1848, 0.2108, 0.1551]],\n",
       "  \n",
       "           [[0.1058, 0.1966, 0.1794, 0.1803, 0.1940, 0.1440],\n",
       "            [0.1243, 0.2108, 0.1993, 0.1819, 0.1685, 0.1153],\n",
       "            [0.1255, 0.1980, 0.2008, 0.1881, 0.1697, 0.1179],\n",
       "            [0.1241, 0.1907, 0.1992, 0.1928, 0.1768, 0.1164],\n",
       "            [0.1173, 0.1843, 0.1894, 0.1927, 0.1938, 0.1225],\n",
       "            [0.1159, 0.1774, 0.1767, 0.1884, 0.2045, 0.1371]],\n",
       "  \n",
       "           [[0.1272, 0.1857, 0.1759, 0.1757, 0.1737, 0.1618],\n",
       "            [0.1044, 0.1654, 0.1610, 0.1713, 0.1782, 0.2196],\n",
       "            [0.1075, 0.1438, 0.1457, 0.1725, 0.1917, 0.2389],\n",
       "            [0.1185, 0.1283, 0.1312, 0.1722, 0.2051, 0.2447],\n",
       "            [0.1189, 0.1447, 0.1362, 0.1722, 0.2095, 0.2185],\n",
       "            [0.1204, 0.1864, 0.1687, 0.1742, 0.1950, 0.1553]],\n",
       "  \n",
       "           [[0.2542, 0.1696, 0.1528, 0.1265, 0.1025, 0.1945],\n",
       "            [0.2145, 0.1804, 0.1726, 0.1499, 0.1274, 0.1552],\n",
       "            [0.1726, 0.1923, 0.1940, 0.1618, 0.1248, 0.1544],\n",
       "            [0.1528, 0.1917, 0.1999, 0.1705, 0.1278, 0.1573],\n",
       "            [0.1446, 0.1793, 0.1956, 0.1792, 0.1342, 0.1671],\n",
       "            [0.1664, 0.1735, 0.1775, 0.1661, 0.1330, 0.1836]],\n",
       "  \n",
       "           [[0.1106, 0.1781, 0.1796, 0.1787, 0.1915, 0.1615],\n",
       "            [0.1318, 0.1930, 0.1721, 0.1596, 0.1672, 0.1764],\n",
       "            [0.1312, 0.1966, 0.1716, 0.1553, 0.1635, 0.1818],\n",
       "            [0.1328, 0.2070, 0.1753, 0.1538, 0.1584, 0.1727],\n",
       "            [0.1373, 0.2182, 0.1809, 0.1561, 0.1549, 0.1526],\n",
       "            [0.1325, 0.2435, 0.2056, 0.1721, 0.1498, 0.0965]],\n",
       "  \n",
       "           [[0.1952, 0.1652, 0.1558, 0.1402, 0.1294, 0.2141],\n",
       "            [0.2243, 0.1447, 0.1305, 0.1209, 0.1286, 0.2509],\n",
       "            [0.2275, 0.1566, 0.1410, 0.1191, 0.1164, 0.2394],\n",
       "            [0.2306, 0.1561, 0.1463, 0.1184, 0.1105, 0.2381],\n",
       "            [0.2370, 0.1345, 0.1332, 0.1182, 0.1202, 0.2568],\n",
       "            [0.2081, 0.1510, 0.1523, 0.1390, 0.1359, 0.2138]]]])])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(TOKENIZER('CCCC', return_tensors='pt')['input_ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([tensor([[[[0.1726, 0.1568, 0.1727, 0.3490, 0.1488],\n",
       "            [0.1110, 0.2015, 0.2282, 0.3161, 0.1432],\n",
       "            [0.1163, 0.2152, 0.2362, 0.2982, 0.1341],\n",
       "            [0.2197, 0.1869, 0.1896, 0.2676, 0.1363],\n",
       "            [0.3595, 0.1261, 0.1005, 0.2566, 0.1573]],\n",
       "  \n",
       "           [[0.0940, 0.3431, 0.3187, 0.0913, 0.1529],\n",
       "            [0.1679, 0.2429, 0.2349, 0.2388, 0.1154],\n",
       "            [0.1389, 0.2090, 0.2227, 0.2471, 0.1822],\n",
       "            [0.0997, 0.2054, 0.1619, 0.2453, 0.2876],\n",
       "            [0.2400, 0.1680, 0.1407, 0.2551, 0.1963]],\n",
       "  \n",
       "           [[0.4902, 0.1760, 0.1867, 0.0225, 0.1246],\n",
       "            [0.1889, 0.1947, 0.1999, 0.1823, 0.2342],\n",
       "            [0.1548, 0.2205, 0.2195, 0.1777, 0.2275],\n",
       "            [0.1706, 0.2176, 0.2285, 0.1408, 0.2425],\n",
       "            [0.0710, 0.2786, 0.2719, 0.1060, 0.2725]],\n",
       "  \n",
       "           [[0.1617, 0.2937, 0.2389, 0.0968, 0.2090],\n",
       "            [0.1281, 0.3208, 0.2768, 0.0876, 0.1868],\n",
       "            [0.1453, 0.3215, 0.2685, 0.0914, 0.1733],\n",
       "            [0.2352, 0.2436, 0.2108, 0.1328, 0.1777],\n",
       "            [0.2451, 0.2384, 0.1999, 0.0482, 0.2685]],\n",
       "  \n",
       "           [[0.3189, 0.0708, 0.1384, 0.2648, 0.2070],\n",
       "            [0.2405, 0.1221, 0.1445, 0.2742, 0.2187],\n",
       "            [0.2332, 0.1332, 0.1516, 0.2743, 0.2077],\n",
       "            [0.3413, 0.0602, 0.0803, 0.1019, 0.4163],\n",
       "            [0.0831, 0.1572, 0.1980, 0.3120, 0.2497]],\n",
       "  \n",
       "           [[0.1171, 0.2606, 0.3001, 0.1062, 0.2160],\n",
       "            [0.2895, 0.1660, 0.1913, 0.0558, 0.2974],\n",
       "            [0.2968, 0.1369, 0.1538, 0.0731, 0.3393],\n",
       "            [0.1294, 0.3019, 0.2676, 0.0745, 0.2266],\n",
       "            [0.2757, 0.2745, 0.3040, 0.0237, 0.1222]],\n",
       "  \n",
       "           [[0.2633, 0.1040, 0.0949, 0.3700, 0.1678],\n",
       "            [0.1360, 0.1923, 0.2051, 0.2977, 0.1690],\n",
       "            [0.1290, 0.1923, 0.2069, 0.2767, 0.1951],\n",
       "            [0.2041, 0.2123, 0.1907, 0.2934, 0.0994],\n",
       "            [0.2438, 0.1804, 0.1834, 0.1828, 0.2096]],\n",
       "  \n",
       "           [[0.0950, 0.1829, 0.2158, 0.2555, 0.2507],\n",
       "            [0.0736, 0.2955, 0.3122, 0.1830, 0.1358],\n",
       "            [0.0666, 0.3082, 0.3291, 0.1701, 0.1260],\n",
       "            [0.1344, 0.2076, 0.2281, 0.2266, 0.2033],\n",
       "            [0.1682, 0.2009, 0.2350, 0.2367, 0.1593]],\n",
       "  \n",
       "           [[0.0677, 0.2815, 0.2850, 0.2340, 0.1319],\n",
       "            [0.1433, 0.2584, 0.2406, 0.2563, 0.1014],\n",
       "            [0.1373, 0.2597, 0.2308, 0.2793, 0.0929],\n",
       "            [0.3114, 0.0796, 0.0675, 0.3245, 0.2169],\n",
       "            [0.2059, 0.2484, 0.2436, 0.1819, 0.1202]],\n",
       "  \n",
       "           [[0.0661, 0.2642, 0.2678, 0.3288, 0.0731],\n",
       "            [0.1235, 0.3030, 0.3037, 0.2498, 0.0200],\n",
       "            [0.0978, 0.3226, 0.3262, 0.2349, 0.0185],\n",
       "            [0.1706, 0.2022, 0.2117, 0.3278, 0.0877],\n",
       "            [0.1759, 0.2736, 0.2893, 0.1783, 0.0829]],\n",
       "  \n",
       "           [[0.1616, 0.0799, 0.0801, 0.4167, 0.2617],\n",
       "            [0.4599, 0.1159, 0.1745, 0.0729, 0.1767],\n",
       "            [0.4516, 0.1052, 0.1581, 0.0946, 0.1905],\n",
       "            [0.1681, 0.1483, 0.1658, 0.2150, 0.3028],\n",
       "            [0.3011, 0.1007, 0.1215, 0.3496, 0.1272]],\n",
       "  \n",
       "           [[0.1874, 0.2122, 0.2101, 0.1311, 0.2592],\n",
       "            [0.0816, 0.3727, 0.3786, 0.0206, 0.1466],\n",
       "            [0.0834, 0.3696, 0.3676, 0.0295, 0.1500],\n",
       "            [0.1914, 0.2758, 0.2917, 0.1218, 0.1192],\n",
       "            [0.2484, 0.2375, 0.2429, 0.1306, 0.1406]]]]),\n",
       "  tensor([[[[0.2564, 0.1663, 0.1835, 0.1985, 0.1953],\n",
       "            [0.1583, 0.1702, 0.1782, 0.3357, 0.1576],\n",
       "            [0.1869, 0.1653, 0.1749, 0.3134, 0.1596],\n",
       "            [0.3045, 0.1035, 0.1245, 0.3336, 0.1338],\n",
       "            [0.3007, 0.2065, 0.1999, 0.1200, 0.1730]],\n",
       "  \n",
       "           [[0.1495, 0.2258, 0.2165, 0.1755, 0.2327],\n",
       "            [0.1027, 0.1999, 0.2275, 0.1781, 0.2917],\n",
       "            [0.0952, 0.2007, 0.2232, 0.1653, 0.3156],\n",
       "            [0.1764, 0.1433, 0.1332, 0.2560, 0.2912],\n",
       "            [0.1476, 0.2383, 0.2611, 0.1198, 0.2332]],\n",
       "  \n",
       "           [[0.2248, 0.1493, 0.1839, 0.2586, 0.1834],\n",
       "            [0.3018, 0.1608, 0.1906, 0.1539, 0.1928],\n",
       "            [0.2655, 0.1428, 0.1837, 0.1904, 0.2176],\n",
       "            [0.2454, 0.2076, 0.2516, 0.1807, 0.1147],\n",
       "            [0.2091, 0.2352, 0.2775, 0.0960, 0.1822]],\n",
       "  \n",
       "           [[0.1908, 0.1812, 0.2148, 0.1335, 0.2797],\n",
       "            [0.1241, 0.1842, 0.1916, 0.3444, 0.1557],\n",
       "            [0.1488, 0.1477, 0.1670, 0.3496, 0.1870],\n",
       "            [0.1322, 0.1931, 0.2206, 0.1664, 0.2878],\n",
       "            [0.1739, 0.1173, 0.1413, 0.3518, 0.2157]],\n",
       "  \n",
       "           [[0.0595, 0.2223, 0.2353, 0.2390, 0.2439],\n",
       "            [0.1117, 0.2669, 0.2755, 0.2133, 0.1327],\n",
       "            [0.0945, 0.2725, 0.2939, 0.2052, 0.1340],\n",
       "            [0.2209, 0.2625, 0.2972, 0.1409, 0.0784],\n",
       "            [0.1084, 0.2894, 0.2552, 0.1993, 0.1477]],\n",
       "  \n",
       "           [[0.2948, 0.1907, 0.2225, 0.0412, 0.2508],\n",
       "            [0.0343, 0.2464, 0.3114, 0.2112, 0.1967],\n",
       "            [0.0486, 0.2445, 0.2968, 0.2276, 0.1825],\n",
       "            [0.1299, 0.2154, 0.1996, 0.2105, 0.2446],\n",
       "            [0.3069, 0.1386, 0.1888, 0.1838, 0.1819]],\n",
       "  \n",
       "           [[0.1763, 0.2622, 0.2684, 0.1765, 0.1166],\n",
       "            [0.1083, 0.2611, 0.2729, 0.1342, 0.2235],\n",
       "            [0.0994, 0.2729, 0.2752, 0.1306, 0.2219],\n",
       "            [0.1404, 0.2582, 0.2461, 0.1327, 0.2224],\n",
       "            [0.0449, 0.2646, 0.2671, 0.1325, 0.2909]],\n",
       "  \n",
       "           [[0.2309, 0.1429, 0.1989, 0.2842, 0.1431],\n",
       "            [0.1117, 0.2369, 0.2270, 0.2260, 0.1984],\n",
       "            [0.1076, 0.2143, 0.2393, 0.2308, 0.2080],\n",
       "            [0.0803, 0.2490, 0.2205, 0.2627, 0.1875],\n",
       "            [0.2447, 0.1137, 0.1784, 0.3199, 0.1433]],\n",
       "  \n",
       "           [[0.1532, 0.2386, 0.2236, 0.2200, 0.1646],\n",
       "            [0.2627, 0.1536, 0.1782, 0.1840, 0.2215],\n",
       "            [0.2537, 0.1531, 0.1621, 0.1816, 0.2495],\n",
       "            [0.2712, 0.1599, 0.1756, 0.2016, 0.1917],\n",
       "            [0.1357, 0.2045, 0.2092, 0.1582, 0.2924]],\n",
       "  \n",
       "           [[0.1027, 0.2602, 0.2759, 0.1808, 0.1804],\n",
       "            [0.0559, 0.2226, 0.2234, 0.1631, 0.3350],\n",
       "            [0.0507, 0.2195, 0.2194, 0.1644, 0.3459],\n",
       "            [0.1749, 0.2047, 0.2198, 0.1735, 0.2271],\n",
       "            [0.1222, 0.1979, 0.1970, 0.3417, 0.1412]],\n",
       "  \n",
       "           [[0.2241, 0.0808, 0.0984, 0.3862, 0.2106],\n",
       "            [0.3249, 0.1335, 0.1512, 0.2557, 0.1347],\n",
       "            [0.3221, 0.1325, 0.1512, 0.2345, 0.1597],\n",
       "            [0.2220, 0.1593, 0.1815, 0.1259, 0.3113],\n",
       "            [0.2488, 0.1247, 0.1461, 0.2593, 0.2210]],\n",
       "  \n",
       "           [[0.2611, 0.1044, 0.0942, 0.2265, 0.3137],\n",
       "            [0.1705, 0.1255, 0.1251, 0.2506, 0.3284],\n",
       "            [0.1581, 0.1433, 0.1229, 0.2613, 0.3144],\n",
       "            [0.2498, 0.1299, 0.0997, 0.2409, 0.2797],\n",
       "            [0.3187, 0.0248, 0.0118, 0.3544, 0.2904]]]]),\n",
       "  tensor([[[[0.2568, 0.2217, 0.2397, 0.1137, 0.1680],\n",
       "            [0.3250, 0.2068, 0.1997, 0.0906, 0.1779],\n",
       "            [0.2990, 0.2304, 0.2169, 0.0946, 0.1591],\n",
       "            [0.2021, 0.1733, 0.2194, 0.2039, 0.2013],\n",
       "            [0.4032, 0.1509, 0.1117, 0.1494, 0.1848]],\n",
       "  \n",
       "           [[0.1315, 0.2028, 0.2012, 0.1762, 0.2883],\n",
       "            [0.1759, 0.1263, 0.1155, 0.1523, 0.4300],\n",
       "            [0.1618, 0.1337, 0.1254, 0.1856, 0.3935],\n",
       "            [0.2540, 0.2206, 0.2188, 0.1581, 0.1486],\n",
       "            [0.1548, 0.1581, 0.1314, 0.2057, 0.3500]],\n",
       "  \n",
       "           [[0.1132, 0.1775, 0.1609, 0.1971, 0.3514],\n",
       "            [0.3139, 0.1966, 0.1348, 0.1432, 0.2114],\n",
       "            [0.3021, 0.2026, 0.1496, 0.1401, 0.2056],\n",
       "            [0.3067, 0.1431, 0.0992, 0.0770, 0.3740],\n",
       "            [0.2592, 0.2012, 0.2445, 0.1565, 0.1386]],\n",
       "  \n",
       "           [[0.0927, 0.2786, 0.2382, 0.2571, 0.1333],\n",
       "            [0.1328, 0.1827, 0.1831, 0.0818, 0.4196],\n",
       "            [0.1467, 0.1434, 0.1394, 0.1184, 0.4522],\n",
       "            [0.2042, 0.2577, 0.1999, 0.1239, 0.2142],\n",
       "            [0.0547, 0.2820, 0.2470, 0.2108, 0.2055]],\n",
       "  \n",
       "           [[0.2297, 0.1075, 0.1159, 0.3822, 0.1647],\n",
       "            [0.2410, 0.1115, 0.1215, 0.2630, 0.2630],\n",
       "            [0.2280, 0.1053, 0.1147, 0.2873, 0.2648],\n",
       "            [0.2670, 0.1438, 0.1689, 0.3388, 0.0814],\n",
       "            [0.1637, 0.1709, 0.2049, 0.3397, 0.1207]],\n",
       "  \n",
       "           [[0.2151, 0.2192, 0.2350, 0.0846, 0.2461],\n",
       "            [0.2800, 0.1297, 0.1476, 0.1489, 0.2939],\n",
       "            [0.2918, 0.1295, 0.1517, 0.1490, 0.2780],\n",
       "            [0.3465, 0.2056, 0.2126, 0.1290, 0.1063],\n",
       "            [0.2524, 0.1811, 0.1861, 0.1647, 0.2157]],\n",
       "  \n",
       "           [[0.1560, 0.2405, 0.3048, 0.1946, 0.1041],\n",
       "            [0.0775, 0.3366, 0.2853, 0.2355, 0.0651],\n",
       "            [0.0880, 0.3452, 0.2967, 0.2058, 0.0642],\n",
       "            [0.2002, 0.2744, 0.2654, 0.1543, 0.1057],\n",
       "            [0.1208, 0.2765, 0.2941, 0.1691, 0.1396]],\n",
       "  \n",
       "           [[0.1544, 0.2597, 0.2415, 0.0899, 0.2545],\n",
       "            [0.2274, 0.1325, 0.1298, 0.1408, 0.3695],\n",
       "            [0.3304, 0.1106, 0.0923, 0.1652, 0.3015],\n",
       "            [0.1414, 0.2685, 0.2689, 0.1361, 0.1851],\n",
       "            [0.3553, 0.1563, 0.1827, 0.1089, 0.1968]],\n",
       "  \n",
       "           [[0.1788, 0.2103, 0.2034, 0.0819, 0.3257],\n",
       "            [0.2276, 0.2951, 0.2698, 0.0861, 0.1215],\n",
       "            [0.2135, 0.3027, 0.2825, 0.0888, 0.1126],\n",
       "            [0.0848, 0.2852, 0.2746, 0.1850, 0.1704],\n",
       "            [0.2189, 0.2839, 0.2770, 0.1186, 0.1015]],\n",
       "  \n",
       "           [[0.3708, 0.1331, 0.1399, 0.1623, 0.1940],\n",
       "            [0.2251, 0.1540, 0.1466, 0.2895, 0.1849],\n",
       "            [0.2020, 0.1502, 0.1451, 0.2884, 0.2144],\n",
       "            [0.2722, 0.1246, 0.0996, 0.2763, 0.2273],\n",
       "            [0.2957, 0.1585, 0.1538, 0.2477, 0.1443]],\n",
       "  \n",
       "           [[0.2539, 0.2237, 0.1277, 0.1837, 0.2111],\n",
       "            [0.3579, 0.1366, 0.1618, 0.2026, 0.1412],\n",
       "            [0.3480, 0.1522, 0.1723, 0.1938, 0.1337],\n",
       "            [0.3520, 0.1770, 0.1572, 0.2286, 0.0851],\n",
       "            [0.3884, 0.1367, 0.1578, 0.1379, 0.1791]],\n",
       "  \n",
       "           [[0.2216, 0.2100, 0.2115, 0.1317, 0.2252],\n",
       "            [0.1949, 0.1531, 0.1655, 0.2916, 0.1949],\n",
       "            [0.1904, 0.1438, 0.1534, 0.2808, 0.2316],\n",
       "            [0.2454, 0.2398, 0.1849, 0.1903, 0.1395],\n",
       "            [0.0773, 0.2932, 0.2337, 0.1328, 0.2630]]]]),\n",
       "  tensor([[[[0.1726, 0.2577, 0.2295, 0.1635, 0.1767],\n",
       "            [0.2497, 0.1829, 0.1585, 0.2846, 0.1243],\n",
       "            [0.2610, 0.1866, 0.1612, 0.2627, 0.1285],\n",
       "            [0.3241, 0.2806, 0.2563, 0.1006, 0.0383],\n",
       "            [0.0685, 0.3753, 0.2960, 0.1595, 0.1007]],\n",
       "  \n",
       "           [[0.3072, 0.2298, 0.2547, 0.1059, 0.1024],\n",
       "            [0.3294, 0.2133, 0.1957, 0.2057, 0.0560],\n",
       "            [0.3533, 0.2191, 0.2023, 0.1881, 0.0373],\n",
       "            [0.1371, 0.2762, 0.2700, 0.1539, 0.1629],\n",
       "            [0.1336, 0.2236, 0.2392, 0.3421, 0.0615]],\n",
       "  \n",
       "           [[0.1372, 0.1347, 0.1149, 0.1660, 0.4472],\n",
       "            [0.1074, 0.1908, 0.1445, 0.1097, 0.4477],\n",
       "            [0.1080, 0.1898, 0.1475, 0.1131, 0.4416],\n",
       "            [0.1912, 0.1961, 0.2169, 0.2990, 0.0967],\n",
       "            [0.0369, 0.2031, 0.1856, 0.2135, 0.3610]],\n",
       "  \n",
       "           [[0.3119, 0.2392, 0.1813, 0.1673, 0.1004],\n",
       "            [0.2369, 0.1241, 0.1259, 0.0787, 0.4343],\n",
       "            [0.2638, 0.0981, 0.1000, 0.0643, 0.4739],\n",
       "            [0.2524, 0.1711, 0.1654, 0.1165, 0.2947],\n",
       "            [0.1291, 0.2570, 0.2165, 0.0882, 0.3092]],\n",
       "  \n",
       "           [[0.2594, 0.0522, 0.0585, 0.2970, 0.3329],\n",
       "            [0.3057, 0.1188, 0.1105, 0.0946, 0.3704],\n",
       "            [0.2703, 0.1226, 0.1091, 0.1226, 0.3755],\n",
       "            [0.2135, 0.0753, 0.0771, 0.3314, 0.3027],\n",
       "            [0.2931, 0.1687, 0.1769, 0.0915, 0.2698]],\n",
       "  \n",
       "           [[0.1513, 0.2233, 0.2332, 0.1287, 0.2636],\n",
       "            [0.1515, 0.2551, 0.2497, 0.1380, 0.2056],\n",
       "            [0.1638, 0.2580, 0.2440, 0.1512, 0.1830],\n",
       "            [0.1755, 0.2323, 0.2044, 0.2461, 0.1416],\n",
       "            [0.2177, 0.2060, 0.1819, 0.2551, 0.1393]],\n",
       "  \n",
       "           [[0.3850, 0.1279, 0.1493, 0.0887, 0.2490],\n",
       "            [0.2841, 0.1953, 0.1860, 0.1199, 0.2147],\n",
       "            [0.2768, 0.2017, 0.1909, 0.1249, 0.2056],\n",
       "            [0.1875, 0.1648, 0.1701, 0.1171, 0.3604],\n",
       "            [0.2060, 0.1832, 0.1948, 0.1396, 0.2764]],\n",
       "  \n",
       "           [[0.2325, 0.2040, 0.1848, 0.1096, 0.2691],\n",
       "            [0.1269, 0.2026, 0.2364, 0.2657, 0.1684],\n",
       "            [0.1259, 0.2048, 0.2465, 0.2480, 0.1748],\n",
       "            [0.1729, 0.2118, 0.2111, 0.1204, 0.2838],\n",
       "            [0.2366, 0.1970, 0.2065, 0.2065, 0.1534]],\n",
       "  \n",
       "           [[0.1597, 0.2886, 0.2631, 0.1064, 0.1822],\n",
       "            [0.2345, 0.1780, 0.1786, 0.2487, 0.1603],\n",
       "            [0.2055, 0.1992, 0.1913, 0.2419, 0.1622],\n",
       "            [0.2496, 0.2170, 0.1832, 0.1733, 0.1769],\n",
       "            [0.1858, 0.2846, 0.2376, 0.1469, 0.1451]],\n",
       "  \n",
       "           [[0.2293, 0.2381, 0.2490, 0.1648, 0.1188],\n",
       "            [0.1531, 0.2414, 0.2480, 0.1890, 0.1685],\n",
       "            [0.1543, 0.2447, 0.2474, 0.1926, 0.1611],\n",
       "            [0.2337, 0.2184, 0.2200, 0.1217, 0.2061],\n",
       "            [0.1474, 0.1519, 0.1360, 0.2896, 0.2751]],\n",
       "  \n",
       "           [[0.1138, 0.2546, 0.3495, 0.1549, 0.1272],\n",
       "            [0.0978, 0.2063, 0.2363, 0.1962, 0.2633],\n",
       "            [0.1032, 0.1976, 0.2305, 0.2061, 0.2626],\n",
       "            [0.0639, 0.1955, 0.1911, 0.1271, 0.4224],\n",
       "            [0.1339, 0.2074, 0.2596, 0.1690, 0.2302]],\n",
       "  \n",
       "           [[0.1581, 0.1052, 0.1370, 0.2792, 0.3206],\n",
       "            [0.1002, 0.0944, 0.1279, 0.2756, 0.4019],\n",
       "            [0.0992, 0.0749, 0.1216, 0.2588, 0.4454],\n",
       "            [0.1232, 0.1657, 0.1984, 0.2203, 0.2923],\n",
       "            [0.1947, 0.1472, 0.1684, 0.2902, 0.1996]]]]),\n",
       "  tensor([[[[0.3190, 0.2026, 0.2171, 0.1472, 0.1141],\n",
       "            [0.2662, 0.1803, 0.1404, 0.1700, 0.2430],\n",
       "            [0.2328, 0.2025, 0.1501, 0.1582, 0.2564],\n",
       "            [0.2148, 0.2543, 0.2563, 0.1136, 0.1611],\n",
       "            [0.2463, 0.2194, 0.1938, 0.1952, 0.1453]],\n",
       "  \n",
       "           [[0.2069, 0.1787, 0.1663, 0.2996, 0.1486],\n",
       "            [0.1711, 0.1470, 0.1506, 0.4122, 0.1192],\n",
       "            [0.1810, 0.1551, 0.1639, 0.3700, 0.1300],\n",
       "            [0.2076, 0.0935, 0.0839, 0.4653, 0.1497],\n",
       "            [0.1611, 0.1450, 0.1293, 0.4312, 0.1334]],\n",
       "  \n",
       "           [[0.1548, 0.1913, 0.2001, 0.2685, 0.1854],\n",
       "            [0.2151, 0.1436, 0.1234, 0.3330, 0.1849],\n",
       "            [0.2065, 0.1434, 0.1261, 0.3408, 0.1832],\n",
       "            [0.1375, 0.2225, 0.2186, 0.2863, 0.1351],\n",
       "            [0.1698, 0.1679, 0.1383, 0.3104, 0.2135]],\n",
       "  \n",
       "           [[0.0940, 0.2710, 0.2838, 0.2832, 0.0680],\n",
       "            [0.1484, 0.3123, 0.3128, 0.1692, 0.0573],\n",
       "            [0.1341, 0.3222, 0.3235, 0.1665, 0.0537],\n",
       "            [0.2304, 0.2135, 0.2411, 0.1664, 0.1485],\n",
       "            [0.2005, 0.2563, 0.2330, 0.1868, 0.1235]],\n",
       "  \n",
       "           [[0.0902, 0.2030, 0.1823, 0.2585, 0.2660],\n",
       "            [0.0958, 0.2062, 0.2572, 0.1706, 0.2702],\n",
       "            [0.0901, 0.1973, 0.2709, 0.1150, 0.3268],\n",
       "            [0.0913, 0.1850, 0.2070, 0.4404, 0.0762],\n",
       "            [0.0786, 0.2640, 0.2359, 0.2697, 0.1518]],\n",
       "  \n",
       "           [[0.2114, 0.1227, 0.1264, 0.3440, 0.1956],\n",
       "            [0.2008, 0.1286, 0.1366, 0.4571, 0.0769],\n",
       "            [0.1699, 0.1497, 0.1588, 0.4546, 0.0670],\n",
       "            [0.2573, 0.1646, 0.1746, 0.2961, 0.1074],\n",
       "            [0.1966, 0.1808, 0.1630, 0.3703, 0.0893]],\n",
       "  \n",
       "           [[0.0742, 0.3490, 0.2747, 0.1204, 0.1818],\n",
       "            [0.0755, 0.2674, 0.2449, 0.2782, 0.1340],\n",
       "            [0.0444, 0.3205, 0.3043, 0.2402, 0.0906],\n",
       "            [0.0581, 0.1597, 0.2158, 0.2138, 0.3527],\n",
       "            [0.1033, 0.2769, 0.2695, 0.1025, 0.2478]],\n",
       "  \n",
       "           [[0.1586, 0.1288, 0.1855, 0.2592, 0.2680],\n",
       "            [0.2007, 0.1233, 0.1515, 0.2921, 0.2324],\n",
       "            [0.1696, 0.1089, 0.1679, 0.3094, 0.2441],\n",
       "            [0.1485, 0.1052, 0.1798, 0.3878, 0.1787],\n",
       "            [0.2087, 0.1376, 0.2292, 0.2805, 0.1440]],\n",
       "  \n",
       "           [[0.4503, 0.1037, 0.1543, 0.1900, 0.1018],\n",
       "            [0.4903, 0.1164, 0.1151, 0.2189, 0.0594],\n",
       "            [0.4859, 0.1119, 0.1149, 0.2036, 0.0836],\n",
       "            [0.4179, 0.1330, 0.1766, 0.1486, 0.1240],\n",
       "            [0.3843, 0.1397, 0.1826, 0.1667, 0.1267]],\n",
       "  \n",
       "           [[0.1819, 0.2132, 0.2434, 0.2489, 0.1127],\n",
       "            [0.0636, 0.2488, 0.2740, 0.1754, 0.2382],\n",
       "            [0.0478, 0.2264, 0.2458, 0.2395, 0.2406],\n",
       "            [0.0956, 0.1872, 0.2590, 0.2893, 0.1689],\n",
       "            [0.1821, 0.1969, 0.2114, 0.2396, 0.1699]],\n",
       "  \n",
       "           [[0.2334, 0.1732, 0.1281, 0.3586, 0.1066],\n",
       "            [0.1140, 0.2506, 0.1968, 0.2663, 0.1722],\n",
       "            [0.0903, 0.2426, 0.1790, 0.2952, 0.1930],\n",
       "            [0.1521, 0.1164, 0.1260, 0.2662, 0.3392],\n",
       "            [0.1011, 0.1912, 0.1942, 0.2197, 0.2938]],\n",
       "  \n",
       "           [[0.3424, 0.1597, 0.1630, 0.0364, 0.2984],\n",
       "            [0.2187, 0.1797, 0.1968, 0.1078, 0.2970],\n",
       "            [0.2271, 0.1924, 0.2093, 0.1136, 0.2577],\n",
       "            [0.1394, 0.2587, 0.2580, 0.1636, 0.1803],\n",
       "            [0.1988, 0.2483, 0.2652, 0.1412, 0.1465]]]]),\n",
       "  tensor([[[[0.3327, 0.1576, 0.1720, 0.2101, 0.1275],\n",
       "            [0.2698, 0.1539, 0.1471, 0.2045, 0.2247],\n",
       "            [0.2662, 0.1694, 0.1577, 0.2065, 0.2002],\n",
       "            [0.2933, 0.1814, 0.1845, 0.2624, 0.0784],\n",
       "            [0.2378, 0.1305, 0.1279, 0.2917, 0.2121]],\n",
       "  \n",
       "           [[0.1408, 0.1610, 0.1685, 0.3129, 0.2169],\n",
       "            [0.1427, 0.2308, 0.2637, 0.1511, 0.2118],\n",
       "            [0.1452, 0.2431, 0.2682, 0.1465, 0.1970],\n",
       "            [0.1114, 0.1837, 0.2315, 0.2568, 0.2167],\n",
       "            [0.1328, 0.1856, 0.2227, 0.2633, 0.1957]],\n",
       "  \n",
       "           [[0.0086, 0.2135, 0.1941, 0.3775, 0.2063],\n",
       "            [0.1595, 0.2391, 0.2256, 0.0997, 0.2761],\n",
       "            [0.1587, 0.2368, 0.2262, 0.1023, 0.2760],\n",
       "            [0.1364, 0.2094, 0.1987, 0.2251, 0.2304],\n",
       "            [0.1699, 0.2288, 0.2427, 0.1511, 0.2075]],\n",
       "  \n",
       "           [[0.1480, 0.2134, 0.2352, 0.2753, 0.1281],\n",
       "            [0.3082, 0.1694, 0.1950, 0.1576, 0.1699],\n",
       "            [0.2939, 0.1711, 0.1973, 0.1611, 0.1766],\n",
       "            [0.0987, 0.2249, 0.2653, 0.2689, 0.1422],\n",
       "            [0.2082, 0.1899, 0.2431, 0.0929, 0.2659]],\n",
       "  \n",
       "           [[0.1272, 0.2002, 0.1949, 0.2233, 0.2543],\n",
       "            [0.0919, 0.1652, 0.1643, 0.2461, 0.3325],\n",
       "            [0.0778, 0.1728, 0.1762, 0.2362, 0.3370],\n",
       "            [0.1685, 0.1657, 0.1766, 0.1410, 0.3482],\n",
       "            [0.1584, 0.1853, 0.1886, 0.2072, 0.2605]],\n",
       "  \n",
       "           [[0.2834, 0.1650, 0.1805, 0.2351, 0.1360],\n",
       "            [0.2239, 0.2484, 0.2630, 0.1399, 0.1248],\n",
       "            [0.2310, 0.2370, 0.2380, 0.1650, 0.1289],\n",
       "            [0.2268, 0.1607, 0.1728, 0.2112, 0.2284],\n",
       "            [0.2548, 0.1672, 0.1668, 0.1933, 0.2179]],\n",
       "  \n",
       "           [[0.1756, 0.2435, 0.2208, 0.1995, 0.1606],\n",
       "            [0.1447, 0.3137, 0.2772, 0.0822, 0.1822],\n",
       "            [0.1740, 0.2698, 0.2522, 0.1079, 0.1960],\n",
       "            [0.1900, 0.2172, 0.1654, 0.3182, 0.1092],\n",
       "            [0.1711, 0.2087, 0.1933, 0.1432, 0.2837]],\n",
       "  \n",
       "           [[0.2155, 0.2853, 0.2477, 0.0938, 0.1577],\n",
       "            [0.1546, 0.3368, 0.3164, 0.0491, 0.1430],\n",
       "            [0.1448, 0.3453, 0.3253, 0.0370, 0.1476],\n",
       "            [0.2639, 0.2660, 0.2203, 0.1822, 0.0676],\n",
       "            [0.1167, 0.2943, 0.3056, 0.0726, 0.2108]],\n",
       "  \n",
       "           [[0.2441, 0.1896, 0.1448, 0.2985, 0.1231],\n",
       "            [0.1334, 0.2202, 0.1400, 0.2983, 0.2080],\n",
       "            [0.1481, 0.2160, 0.1398, 0.2857, 0.2104],\n",
       "            [0.2711, 0.2379, 0.2568, 0.1473, 0.0868],\n",
       "            [0.2598, 0.1735, 0.1400, 0.2951, 0.1316]],\n",
       "  \n",
       "           [[0.2921, 0.1547, 0.1554, 0.2346, 0.1632],\n",
       "            [0.3620, 0.1456, 0.1471, 0.1812, 0.1640],\n",
       "            [0.3480, 0.1357, 0.1412, 0.1896, 0.1855],\n",
       "            [0.2627, 0.0946, 0.1103, 0.2223, 0.3100],\n",
       "            [0.3330, 0.1367, 0.1329, 0.1943, 0.2031]],\n",
       "  \n",
       "           [[0.4130, 0.1283, 0.1307, 0.1285, 0.1994],\n",
       "            [0.4265, 0.0724, 0.0701, 0.0904, 0.3406],\n",
       "            [0.4029, 0.0962, 0.0879, 0.0965, 0.3164],\n",
       "            [0.4320, 0.1111, 0.1173, 0.1558, 0.1837],\n",
       "            [0.4117, 0.1450, 0.1065, 0.1635, 0.1732]],\n",
       "  \n",
       "           [[0.2481, 0.3120, 0.3608, 0.0260, 0.0530],\n",
       "            [0.2208, 0.3422, 0.2928, 0.1247, 0.0195],\n",
       "            [0.2073, 0.3393, 0.3177, 0.1071, 0.0286],\n",
       "            [0.1436, 0.3053, 0.2947, 0.1174, 0.1390],\n",
       "            [0.2405, 0.2408, 0.2462, 0.1309, 0.1416]]]]),\n",
       "  tensor([[[[0.1441, 0.1432, 0.1096, 0.2881, 0.3150],\n",
       "            [0.1661, 0.2327, 0.2433, 0.2726, 0.0854],\n",
       "            [0.1372, 0.2393, 0.2492, 0.2728, 0.1015],\n",
       "            [0.2145, 0.1314, 0.1533, 0.2494, 0.2514],\n",
       "            [0.2158, 0.1084, 0.1537, 0.0402, 0.4819]],\n",
       "  \n",
       "           [[0.1873, 0.2410, 0.2079, 0.0906, 0.2732],\n",
       "            [0.2122, 0.1845, 0.1871, 0.1892, 0.2269],\n",
       "            [0.2140, 0.1903, 0.1810, 0.1626, 0.2521],\n",
       "            [0.1758, 0.1960, 0.1976, 0.1401, 0.2905],\n",
       "            [0.2849, 0.1686, 0.1641, 0.1709, 0.2114]],\n",
       "  \n",
       "           [[0.0684, 0.1693, 0.1685, 0.3657, 0.2281],\n",
       "            [0.1650, 0.1446, 0.1487, 0.3397, 0.2021],\n",
       "            [0.1557, 0.1589, 0.1469, 0.3337, 0.2048],\n",
       "            [0.1006, 0.2232, 0.1821, 0.2401, 0.2540],\n",
       "            [0.1695, 0.1626, 0.1646, 0.3206, 0.1828]],\n",
       "  \n",
       "           [[0.0815, 0.1800, 0.1514, 0.1714, 0.4156],\n",
       "            [0.0929, 0.2346, 0.2175, 0.0396, 0.4154],\n",
       "            [0.1277, 0.2609, 0.2447, 0.0354, 0.3313],\n",
       "            [0.1444, 0.2574, 0.2212, 0.1293, 0.2477],\n",
       "            [0.1728, 0.2439, 0.1888, 0.0897, 0.3048]],\n",
       "  \n",
       "           [[0.2098, 0.1869, 0.2217, 0.2142, 0.1675],\n",
       "            [0.2382, 0.1563, 0.2021, 0.1937, 0.2098],\n",
       "            [0.2012, 0.1676, 0.2227, 0.1952, 0.2132],\n",
       "            [0.1537, 0.1890, 0.2322, 0.2137, 0.2114],\n",
       "            [0.1451, 0.1578, 0.2082, 0.2496, 0.2393]],\n",
       "  \n",
       "           [[0.1567, 0.2368, 0.2225, 0.1416, 0.2423],\n",
       "            [0.0618, 0.0922, 0.1302, 0.2267, 0.4891],\n",
       "            [0.0681, 0.1105, 0.1496, 0.2020, 0.4698],\n",
       "            [0.0781, 0.1501, 0.1816, 0.1809, 0.4092],\n",
       "            [0.0481, 0.1405, 0.1959, 0.1625, 0.4529]],\n",
       "  \n",
       "           [[0.1744, 0.2080, 0.2140, 0.2972, 0.1064],\n",
       "            [0.1597, 0.1623, 0.1682, 0.3049, 0.2049],\n",
       "            [0.1620, 0.1696, 0.1790, 0.2924, 0.1970],\n",
       "            [0.1704, 0.2434, 0.2769, 0.1571, 0.1522],\n",
       "            [0.2283, 0.2105, 0.2349, 0.1922, 0.1341]],\n",
       "  \n",
       "           [[0.1628, 0.2729, 0.2547, 0.1626, 0.1470],\n",
       "            [0.1550, 0.2477, 0.2391, 0.1535, 0.2046],\n",
       "            [0.1776, 0.2126, 0.2160, 0.1784, 0.2155],\n",
       "            [0.1326, 0.2507, 0.2320, 0.1782, 0.2065],\n",
       "            [0.1855, 0.1736, 0.1381, 0.2677, 0.2351]],\n",
       "  \n",
       "           [[0.3473, 0.1983, 0.1843, 0.1748, 0.0953],\n",
       "            [0.2752, 0.2679, 0.2470, 0.1353, 0.0745],\n",
       "            [0.2879, 0.2604, 0.2395, 0.1395, 0.0727],\n",
       "            [0.2828, 0.2288, 0.2173, 0.1192, 0.1520],\n",
       "            [0.2349, 0.2190, 0.1889, 0.2064, 0.1509]],\n",
       "  \n",
       "           [[0.1083, 0.2854, 0.2967, 0.1429, 0.1666],\n",
       "            [0.2088, 0.1967, 0.2227, 0.2729, 0.0988],\n",
       "            [0.2099, 0.1855, 0.2063, 0.3015, 0.0968],\n",
       "            [0.1363, 0.2568, 0.2764, 0.1875, 0.1430],\n",
       "            [0.0761, 0.2364, 0.2622, 0.0316, 0.3938]],\n",
       "  \n",
       "           [[0.2366, 0.2845, 0.2985, 0.0384, 0.1422],\n",
       "            [0.2940, 0.2036, 0.1876, 0.2147, 0.1000],\n",
       "            [0.2836, 0.1986, 0.1852, 0.2220, 0.1105],\n",
       "            [0.2611, 0.1235, 0.1392, 0.3037, 0.1725],\n",
       "            [0.2437, 0.1794, 0.1826, 0.2954, 0.0989]],\n",
       "  \n",
       "           [[0.1114, 0.3043, 0.2821, 0.0829, 0.2193],\n",
       "            [0.1856, 0.2197, 0.2438, 0.1382, 0.2127],\n",
       "            [0.1892, 0.1909, 0.2236, 0.1484, 0.2478],\n",
       "            [0.1588, 0.2555, 0.2402, 0.0587, 0.2867],\n",
       "            [0.2067, 0.2202, 0.2185, 0.1020, 0.2526]]]]),\n",
       "  tensor([[[[0.1694, 0.0435, 0.0477, 0.3501, 0.3893],\n",
       "            [0.1021, 0.0985, 0.0947, 0.4630, 0.2415],\n",
       "            [0.1109, 0.0951, 0.0845, 0.4526, 0.2568],\n",
       "            [0.1297, 0.1179, 0.1268, 0.3354, 0.2902],\n",
       "            [0.1797, 0.1242, 0.0882, 0.2710, 0.3369]],\n",
       "  \n",
       "           [[0.3723, 0.2161, 0.1598, 0.1172, 0.1346],\n",
       "            [0.2950, 0.2702, 0.2145, 0.0817, 0.1386],\n",
       "            [0.3568, 0.2336, 0.1877, 0.0937, 0.1283],\n",
       "            [0.2372, 0.2769, 0.2217, 0.1111, 0.1532],\n",
       "            [0.2682, 0.2566, 0.2178, 0.1004, 0.1571]],\n",
       "  \n",
       "           [[0.1063, 0.2256, 0.2367, 0.1494, 0.2820],\n",
       "            [0.1910, 0.1434, 0.1802, 0.2278, 0.2575],\n",
       "            [0.2067, 0.1503, 0.1845, 0.2085, 0.2500],\n",
       "            [0.2065, 0.2131, 0.2071, 0.1563, 0.2171],\n",
       "            [0.1162, 0.1791, 0.1774, 0.2036, 0.3238]],\n",
       "  \n",
       "           [[0.2415, 0.1699, 0.1564, 0.2126, 0.2196],\n",
       "            [0.2371, 0.1652, 0.1728, 0.2185, 0.2064],\n",
       "            [0.2237, 0.1633, 0.1754, 0.2195, 0.2181],\n",
       "            [0.2717, 0.2201, 0.1764, 0.2075, 0.1242],\n",
       "            [0.1555, 0.1644, 0.2059, 0.2593, 0.2149]],\n",
       "  \n",
       "           [[0.3125, 0.1046, 0.1418, 0.1849, 0.2562],\n",
       "            [0.1830, 0.0468, 0.0949, 0.3142, 0.3612],\n",
       "            [0.1743, 0.0649, 0.1075, 0.3134, 0.3400],\n",
       "            [0.3538, 0.1146, 0.1255, 0.1601, 0.2461],\n",
       "            [0.2880, 0.1101, 0.1158, 0.2319, 0.2542]],\n",
       "  \n",
       "           [[0.2418, 0.1961, 0.2146, 0.2739, 0.0736],\n",
       "            [0.1490, 0.2211, 0.2362, 0.2872, 0.1066],\n",
       "            [0.1296, 0.2268, 0.2345, 0.2796, 0.1296],\n",
       "            [0.4046, 0.1924, 0.1067, 0.1659, 0.1304],\n",
       "            [0.0678, 0.2182, 0.2270, 0.3507, 0.1363]],\n",
       "  \n",
       "           [[0.1175, 0.1797, 0.1390, 0.4309, 0.1330],\n",
       "            [0.0820, 0.0834, 0.0687, 0.5107, 0.2551],\n",
       "            [0.0771, 0.1026, 0.0796, 0.4590, 0.2817],\n",
       "            [0.1864, 0.0813, 0.0690, 0.5593, 0.1040],\n",
       "            [0.0656, 0.2245, 0.2016, 0.3354, 0.1728]],\n",
       "  \n",
       "           [[0.2135, 0.1725, 0.1742, 0.1942, 0.2456],\n",
       "            [0.2039, 0.1388, 0.1575, 0.2081, 0.2917],\n",
       "            [0.2068, 0.1437, 0.1624, 0.1949, 0.2922],\n",
       "            [0.2396, 0.1101, 0.1342, 0.1769, 0.3391],\n",
       "            [0.2431, 0.1512, 0.1494, 0.1648, 0.2915]],\n",
       "  \n",
       "           [[0.2588, 0.1313, 0.1139, 0.3606, 0.1355],\n",
       "            [0.1894, 0.1769, 0.1944, 0.3058, 0.1336],\n",
       "            [0.2041, 0.1354, 0.1642, 0.3335, 0.1628],\n",
       "            [0.2353, 0.1832, 0.1770, 0.2731, 0.1314],\n",
       "            [0.2810, 0.1370, 0.1375, 0.2560, 0.1885]],\n",
       "  \n",
       "           [[0.2632, 0.1484, 0.1505, 0.1500, 0.2880],\n",
       "            [0.3131, 0.1549, 0.1617, 0.1295, 0.2407],\n",
       "            [0.3391, 0.1373, 0.1496, 0.1373, 0.2368],\n",
       "            [0.2781, 0.1603, 0.1759, 0.2342, 0.1516],\n",
       "            [0.4049, 0.0312, 0.0549, 0.1648, 0.3442]],\n",
       "  \n",
       "           [[0.1588, 0.2486, 0.2787, 0.2327, 0.0813],\n",
       "            [0.4886, 0.0066, 0.0522, 0.3883, 0.0644],\n",
       "            [0.4936, 0.0148, 0.0541, 0.3660, 0.0715],\n",
       "            [0.1989, 0.1277, 0.1580, 0.3438, 0.1716],\n",
       "            [0.1134, 0.2574, 0.2852, 0.2347, 0.1093]],\n",
       "  \n",
       "           [[0.2777, 0.1709, 0.1653, 0.1745, 0.2115],\n",
       "            [0.1960, 0.2578, 0.2494, 0.1750, 0.1218],\n",
       "            [0.1875, 0.2697, 0.2649, 0.1605, 0.1174],\n",
       "            [0.2193, 0.2003, 0.2039, 0.2026, 0.1738],\n",
       "            [0.1897, 0.1738, 0.1660, 0.2153, 0.2552]]]]),\n",
       "  tensor([[[[0.2022, 0.2639, 0.2623, 0.1950, 0.0765],\n",
       "            [0.1605, 0.2516, 0.2581, 0.2534, 0.0764],\n",
       "            [0.1493, 0.2497, 0.2749, 0.2359, 0.0902],\n",
       "            [0.2063, 0.2348, 0.2450, 0.2066, 0.1074],\n",
       "            [0.1972, 0.2577, 0.2589, 0.1671, 0.1191]],\n",
       "  \n",
       "           [[0.3460, 0.1994, 0.2133, 0.1375, 0.1038],\n",
       "            [0.3436, 0.1559, 0.1596, 0.2889, 0.0520],\n",
       "            [0.4041, 0.1308, 0.1341, 0.2852, 0.0457],\n",
       "            [0.2457, 0.2313, 0.2286, 0.2543, 0.0402],\n",
       "            [0.1787, 0.1730, 0.1591, 0.2612, 0.2280]],\n",
       "  \n",
       "           [[0.2278, 0.2027, 0.2477, 0.1873, 0.1345],\n",
       "            [0.2185, 0.1357, 0.2079, 0.2128, 0.2250],\n",
       "            [0.2132, 0.1254, 0.1890, 0.1944, 0.2779],\n",
       "            [0.2355, 0.1991, 0.2192, 0.1259, 0.2203],\n",
       "            [0.2536, 0.1666, 0.2139, 0.0630, 0.3029]],\n",
       "  \n",
       "           [[0.1247, 0.3633, 0.3603, 0.0992, 0.0526],\n",
       "            [0.1672, 0.2605, 0.2253, 0.1354, 0.2116],\n",
       "            [0.1952, 0.2775, 0.2257, 0.1361, 0.1655],\n",
       "            [0.1849, 0.2741, 0.2526, 0.1793, 0.1091],\n",
       "            [0.1560, 0.3239, 0.2829, 0.1280, 0.1091]],\n",
       "  \n",
       "           [[0.3102, 0.1337, 0.1210, 0.2418, 0.1933],\n",
       "            [0.1793, 0.2014, 0.1910, 0.2493, 0.1789],\n",
       "            [0.1562, 0.2071, 0.1932, 0.2764, 0.1671],\n",
       "            [0.3021, 0.1781, 0.1650, 0.2136, 0.1412],\n",
       "            [0.2741, 0.1275, 0.1246, 0.2784, 0.1954]],\n",
       "  \n",
       "           [[0.1723, 0.2399, 0.2193, 0.2177, 0.1507],\n",
       "            [0.1014, 0.2576, 0.2451, 0.1691, 0.2269],\n",
       "            [0.1344, 0.2267, 0.2312, 0.1718, 0.2360],\n",
       "            [0.2102, 0.2338, 0.2126, 0.1951, 0.1484],\n",
       "            [0.0549, 0.2937, 0.2741, 0.1721, 0.2053]],\n",
       "  \n",
       "           [[0.1834, 0.2401, 0.2473, 0.1560, 0.1732],\n",
       "            [0.2549, 0.1617, 0.1772, 0.2185, 0.1877],\n",
       "            [0.2568, 0.1616, 0.1769, 0.2236, 0.1810],\n",
       "            [0.2190, 0.1550, 0.1425, 0.2485, 0.2350],\n",
       "            [0.1996, 0.2035, 0.1862, 0.1819, 0.2289]],\n",
       "  \n",
       "           [[0.1700, 0.2075, 0.2014, 0.2016, 0.2196],\n",
       "            [0.2010, 0.1598, 0.1984, 0.1882, 0.2525],\n",
       "            [0.2086, 0.1573, 0.1890, 0.1968, 0.2482],\n",
       "            [0.1792, 0.1979, 0.1872, 0.2088, 0.2271],\n",
       "            [0.2361, 0.1629, 0.1546, 0.2283, 0.2181]],\n",
       "  \n",
       "           [[0.1309, 0.1491, 0.1490, 0.2252, 0.3459],\n",
       "            [0.0377, 0.2264, 0.1909, 0.1159, 0.4291],\n",
       "            [0.0701, 0.2313, 0.1898, 0.1087, 0.4002],\n",
       "            [0.1124, 0.2326, 0.2330, 0.0869, 0.3351],\n",
       "            [0.1550, 0.2034, 0.1797, 0.2324, 0.2294]],\n",
       "  \n",
       "           [[0.1147, 0.1461, 0.1876, 0.3020, 0.2496],\n",
       "            [0.1558, 0.1972, 0.1884, 0.2176, 0.2411],\n",
       "            [0.1625, 0.2067, 0.1907, 0.2037, 0.2365],\n",
       "            [0.1763, 0.1316, 0.1206, 0.2757, 0.2958],\n",
       "            [0.1824, 0.1603, 0.1561, 0.2577, 0.2435]],\n",
       "  \n",
       "           [[0.2498, 0.2161, 0.2280, 0.1195, 0.1866],\n",
       "            [0.2484, 0.2076, 0.2268, 0.1787, 0.1386],\n",
       "            [0.2442, 0.2087, 0.2216, 0.1760, 0.1495],\n",
       "            [0.3727, 0.2182, 0.2078, 0.0797, 0.1215],\n",
       "            [0.3284, 0.2142, 0.2207, 0.0895, 0.1472]],\n",
       "  \n",
       "           [[0.3083, 0.1554, 0.1606, 0.2197, 0.1560],\n",
       "            [0.1800, 0.1833, 0.2254, 0.1703, 0.2410],\n",
       "            [0.1532, 0.1875, 0.2323, 0.1934, 0.2337],\n",
       "            [0.1891, 0.1441, 0.1703, 0.2590, 0.2374],\n",
       "            [0.1414, 0.2450, 0.2580, 0.1222, 0.2335]]]]),\n",
       "  tensor([[[[0.2554, 0.2378, 0.2286, 0.1315, 0.1466],\n",
       "            [0.2621, 0.2301, 0.2282, 0.1576, 0.1220],\n",
       "            [0.2410, 0.2277, 0.2103, 0.1848, 0.1362],\n",
       "            [0.1820, 0.2338, 0.2401, 0.1936, 0.1504],\n",
       "            [0.1980, 0.2351, 0.2813, 0.1904, 0.0952]],\n",
       "  \n",
       "           [[0.2514, 0.2806, 0.2243, 0.0586, 0.1851],\n",
       "            [0.2961, 0.2583, 0.2020, 0.0854, 0.1581],\n",
       "            [0.3029, 0.2455, 0.2084, 0.1003, 0.1429],\n",
       "            [0.3742, 0.2539, 0.1911, 0.0746, 0.1061],\n",
       "            [0.2526, 0.2339, 0.2320, 0.0448, 0.2366]],\n",
       "  \n",
       "           [[0.1418, 0.1999, 0.1824, 0.1276, 0.3484],\n",
       "            [0.2491, 0.2134, 0.1524, 0.1101, 0.2751],\n",
       "            [0.2974, 0.1926, 0.1360, 0.1214, 0.2526],\n",
       "            [0.1333, 0.1447, 0.1784, 0.1954, 0.3481],\n",
       "            [0.2561, 0.1516, 0.1582, 0.1384, 0.2957]],\n",
       "  \n",
       "           [[0.1787, 0.1397, 0.2049, 0.3208, 0.1559],\n",
       "            [0.1991, 0.1299, 0.1762, 0.2887, 0.2061],\n",
       "            [0.2109, 0.1306, 0.1740, 0.2936, 0.1909],\n",
       "            [0.2641, 0.1333, 0.1752, 0.2500, 0.1774],\n",
       "            [0.2520, 0.1280, 0.1709, 0.2851, 0.1640]],\n",
       "  \n",
       "           [[0.2506, 0.1351, 0.1224, 0.2719, 0.2200],\n",
       "            [0.1726, 0.1762, 0.1659, 0.2842, 0.2011],\n",
       "            [0.1629, 0.1905, 0.1791, 0.2807, 0.1868],\n",
       "            [0.2670, 0.1543, 0.1385, 0.2205, 0.2197],\n",
       "            [0.1918, 0.1863, 0.1711, 0.2501, 0.2007]],\n",
       "  \n",
       "           [[0.3073, 0.1476, 0.1141, 0.2612, 0.1697],\n",
       "            [0.1960, 0.1963, 0.1764, 0.2152, 0.2161],\n",
       "            [0.1493, 0.2141, 0.2115, 0.1814, 0.2437],\n",
       "            [0.2428, 0.1850, 0.1491, 0.1654, 0.2577],\n",
       "            [0.2015, 0.2386, 0.2270, 0.1383, 0.1945]],\n",
       "  \n",
       "           [[0.2137, 0.1777, 0.1578, 0.2081, 0.2428],\n",
       "            [0.1798, 0.1991, 0.2143, 0.2737, 0.1331],\n",
       "            [0.1721, 0.1986, 0.2008, 0.2851, 0.1434],\n",
       "            [0.1825, 0.1827, 0.1022, 0.2602, 0.2724],\n",
       "            [0.1849, 0.2241, 0.1830, 0.2058, 0.2022]],\n",
       "  \n",
       "           [[0.2284, 0.2486, 0.2329, 0.0809, 0.2092],\n",
       "            [0.1865, 0.3042, 0.3043, 0.0818, 0.1232],\n",
       "            [0.1962, 0.2734, 0.2645, 0.1194, 0.1464],\n",
       "            [0.1815, 0.2167, 0.2208, 0.1987, 0.1822],\n",
       "            [0.1023, 0.2431, 0.2514, 0.1810, 0.2222]],\n",
       "  \n",
       "           [[0.1032, 0.1948, 0.2153, 0.2117, 0.2750],\n",
       "            [0.1506, 0.1403, 0.2001, 0.2831, 0.2259],\n",
       "            [0.1772, 0.1867, 0.2197, 0.2290, 0.1875],\n",
       "            [0.1639, 0.2016, 0.1828, 0.2302, 0.2215],\n",
       "            [0.1702, 0.2377, 0.2142, 0.2033, 0.1746]],\n",
       "  \n",
       "           [[0.1901, 0.2646, 0.2524, 0.1038, 0.1891],\n",
       "            [0.1961, 0.2703, 0.2228, 0.1526, 0.1582],\n",
       "            [0.2165, 0.2420, 0.1793, 0.1560, 0.2062],\n",
       "            [0.2162, 0.2623, 0.2099, 0.0861, 0.2255],\n",
       "            [0.2235, 0.2392, 0.2034, 0.1102, 0.2237]],\n",
       "  \n",
       "           [[0.1668, 0.2130, 0.2493, 0.1881, 0.1828],\n",
       "            [0.2371, 0.2099, 0.2327, 0.1549, 0.1654],\n",
       "            [0.2836, 0.1797, 0.2230, 0.1367, 0.1770],\n",
       "            [0.3749, 0.1344, 0.1466, 0.1478, 0.1963],\n",
       "            [0.2565, 0.1605, 0.1870, 0.2142, 0.1818]],\n",
       "  \n",
       "           [[0.2477, 0.1844, 0.1575, 0.3198, 0.0906],\n",
       "            [0.2663, 0.1558, 0.1666, 0.3054, 0.1059],\n",
       "            [0.2746, 0.1530, 0.1646, 0.3026, 0.1052],\n",
       "            [0.1857, 0.2368, 0.1971, 0.2743, 0.1061],\n",
       "            [0.1987, 0.1766, 0.1814, 0.3321, 0.1113]]]]),\n",
       "  tensor([[[[0.2511, 0.2297, 0.2299, 0.1205, 0.1688],\n",
       "            [0.2351, 0.2089, 0.2089, 0.1489, 0.1983],\n",
       "            [0.2408, 0.2078, 0.1990, 0.1407, 0.2116],\n",
       "            [0.2174, 0.2334, 0.2278, 0.1567, 0.1647],\n",
       "            [0.2337, 0.2663, 0.2580, 0.0656, 0.1764]],\n",
       "  \n",
       "           [[0.1654, 0.1737, 0.1770, 0.3000, 0.1839],\n",
       "            [0.0724, 0.2914, 0.2941, 0.1941, 0.1479],\n",
       "            [0.0867, 0.2914, 0.3021, 0.1653, 0.1545],\n",
       "            [0.1272, 0.2394, 0.2470, 0.2374, 0.1490],\n",
       "            [0.1939, 0.2041, 0.2124, 0.2819, 0.1077]],\n",
       "  \n",
       "           [[0.1725, 0.1844, 0.1662, 0.2275, 0.2494],\n",
       "            [0.1384, 0.1725, 0.1938, 0.2047, 0.2907],\n",
       "            [0.1217, 0.1622, 0.1881, 0.2084, 0.3197],\n",
       "            [0.1412, 0.1824, 0.1835, 0.2593, 0.2335],\n",
       "            [0.1733, 0.1853, 0.1594, 0.2261, 0.2559]],\n",
       "  \n",
       "           [[0.2099, 0.1828, 0.1994, 0.1275, 0.2804],\n",
       "            [0.1832, 0.2823, 0.2522, 0.0544, 0.2279],\n",
       "            [0.1894, 0.2481, 0.2253, 0.0628, 0.2742],\n",
       "            [0.2486, 0.2253, 0.1990, 0.0946, 0.2325],\n",
       "            [0.2382, 0.2375, 0.1826, 0.0892, 0.2524]],\n",
       "  \n",
       "           [[0.0544, 0.2213, 0.2476, 0.1863, 0.2904],\n",
       "            [0.0840, 0.1617, 0.1992, 0.2355, 0.3196],\n",
       "            [0.0976, 0.1584, 0.1997, 0.2568, 0.2875],\n",
       "            [0.1494, 0.1614, 0.1688, 0.2372, 0.2832],\n",
       "            [0.1362, 0.2052, 0.2179, 0.2015, 0.2392]],\n",
       "  \n",
       "           [[0.2539, 0.1713, 0.1614, 0.2431, 0.1702],\n",
       "            [0.1914, 0.2037, 0.2274, 0.1897, 0.1879],\n",
       "            [0.1592, 0.2181, 0.2343, 0.1884, 0.2000],\n",
       "            [0.1732, 0.2413, 0.2430, 0.1832, 0.1593],\n",
       "            [0.1664, 0.2497, 0.2303, 0.1361, 0.2175]],\n",
       "  \n",
       "           [[0.1110, 0.1795, 0.2126, 0.2222, 0.2747],\n",
       "            [0.1512, 0.2643, 0.2863, 0.1482, 0.1500],\n",
       "            [0.2050, 0.2690, 0.2737, 0.1347, 0.1177],\n",
       "            [0.2582, 0.2638, 0.2393, 0.1614, 0.0773],\n",
       "            [0.1643, 0.2876, 0.2549, 0.1366, 0.1567]],\n",
       "  \n",
       "           [[0.3234, 0.1610, 0.0891, 0.1492, 0.2773],\n",
       "            [0.2896, 0.1886, 0.1521, 0.1027, 0.2670],\n",
       "            [0.2534, 0.2178, 0.1922, 0.0856, 0.2509],\n",
       "            [0.2646, 0.2030, 0.1949, 0.0677, 0.2697],\n",
       "            [0.2448, 0.2415, 0.2240, 0.0647, 0.2250]],\n",
       "  \n",
       "           [[0.1493, 0.2476, 0.2517, 0.1759, 0.1756],\n",
       "            [0.1275, 0.2640, 0.2822, 0.1791, 0.1473],\n",
       "            [0.1052, 0.2571, 0.2724, 0.1661, 0.1992],\n",
       "            [0.1211, 0.2267, 0.2124, 0.2238, 0.2161],\n",
       "            [0.1822, 0.2384, 0.2189, 0.2114, 0.1491]],\n",
       "  \n",
       "           [[0.2517, 0.1961, 0.1906, 0.1660, 0.1956],\n",
       "            [0.1556, 0.3079, 0.2719, 0.1399, 0.1248],\n",
       "            [0.1789, 0.2840, 0.2401, 0.1300, 0.1670],\n",
       "            [0.2203, 0.2419, 0.2151, 0.1419, 0.1807],\n",
       "            [0.2542, 0.1590, 0.1473, 0.1826, 0.2570]],\n",
       "  \n",
       "           [[0.2355, 0.1851, 0.1789, 0.1870, 0.2134],\n",
       "            [0.2358, 0.1897, 0.2012, 0.1068, 0.2665],\n",
       "            [0.2435, 0.1970, 0.2103, 0.1098, 0.2393],\n",
       "            [0.2666, 0.2085, 0.2083, 0.0946, 0.2220],\n",
       "            [0.1713, 0.2470, 0.2295, 0.0911, 0.2612]],\n",
       "  \n",
       "           [[0.2419, 0.1357, 0.1274, 0.2240, 0.2711],\n",
       "            [0.2048, 0.1275, 0.1376, 0.2619, 0.2683],\n",
       "            [0.1937, 0.1344, 0.1413, 0.2641, 0.2665],\n",
       "            [0.1736, 0.1500, 0.1595, 0.2258, 0.2911],\n",
       "            [0.2602, 0.1666, 0.1594, 0.1450, 0.2688]]]]),\n",
       "  tensor([[[[0.1729, 0.2161, 0.2109, 0.1111, 0.2889],\n",
       "            [0.1364, 0.2463, 0.2371, 0.1550, 0.2254],\n",
       "            [0.1260, 0.2419, 0.2457, 0.1590, 0.2274],\n",
       "            [0.1616, 0.2421, 0.2474, 0.0888, 0.2601],\n",
       "            [0.1500, 0.2206, 0.2152, 0.1112, 0.3030]],\n",
       "  \n",
       "           [[0.2361, 0.1573, 0.1486, 0.2416, 0.2165],\n",
       "            [0.2766, 0.1787, 0.1488, 0.1741, 0.2218],\n",
       "            [0.2782, 0.1892, 0.1558, 0.1674, 0.2093],\n",
       "            [0.2780, 0.1915, 0.1617, 0.1902, 0.1786],\n",
       "            [0.2856, 0.1652, 0.1422, 0.1864, 0.2207]],\n",
       "  \n",
       "           [[0.2226, 0.1349, 0.1261, 0.2418, 0.2747],\n",
       "            [0.1604, 0.1576, 0.1657, 0.2098, 0.3065],\n",
       "            [0.1630, 0.1869, 0.1749, 0.2165, 0.2588],\n",
       "            [0.1556, 0.2362, 0.1975, 0.1881, 0.2226],\n",
       "            [0.0651, 0.2481, 0.2641, 0.1360, 0.2868]],\n",
       "  \n",
       "           [[0.2374, 0.2828, 0.2494, 0.1161, 0.1143],\n",
       "            [0.2349, 0.2827, 0.2418, 0.1266, 0.1140],\n",
       "            [0.2298, 0.2786, 0.2393, 0.1266, 0.1257],\n",
       "            [0.1888, 0.2660, 0.2428, 0.1501, 0.1524],\n",
       "            [0.2046, 0.2589, 0.2219, 0.1377, 0.1769]],\n",
       "  \n",
       "           [[0.2465, 0.2136, 0.1943, 0.0850, 0.2606],\n",
       "            [0.2483, 0.2274, 0.2001, 0.1026, 0.2216],\n",
       "            [0.2604, 0.2345, 0.2013, 0.0647, 0.2389],\n",
       "            [0.3146, 0.2015, 0.1794, 0.0766, 0.2279],\n",
       "            [0.1994, 0.2544, 0.2242, 0.0755, 0.2465]],\n",
       "  \n",
       "           [[0.2742, 0.2231, 0.1937, 0.2040, 0.1050],\n",
       "            [0.2384, 0.2138, 0.1995, 0.2103, 0.1381],\n",
       "            [0.1742, 0.1874, 0.2032, 0.2473, 0.1878],\n",
       "            [0.1515, 0.1579, 0.2050, 0.3125, 0.1730],\n",
       "            [0.2366, 0.0460, 0.0797, 0.3648, 0.2729]],\n",
       "  \n",
       "           [[0.2149, 0.2137, 0.2430, 0.1713, 0.1570],\n",
       "            [0.2420, 0.2087, 0.2407, 0.1723, 0.1363],\n",
       "            [0.2713, 0.2063, 0.2304, 0.1736, 0.1185],\n",
       "            [0.2576, 0.2278, 0.2336, 0.1660, 0.1150],\n",
       "            [0.2003, 0.2568, 0.2323, 0.1687, 0.1420]],\n",
       "  \n",
       "           [[0.1254, 0.2653, 0.2430, 0.1620, 0.2043],\n",
       "            [0.1346, 0.2717, 0.2662, 0.1598, 0.1677],\n",
       "            [0.1405, 0.2541, 0.2661, 0.1601, 0.1792],\n",
       "            [0.1488, 0.2353, 0.2515, 0.1858, 0.1785],\n",
       "            [0.1383, 0.2264, 0.2419, 0.1949, 0.1984]],\n",
       "  \n",
       "           [[0.1721, 0.2519, 0.2354, 0.1443, 0.1964],\n",
       "            [0.1585, 0.2163, 0.2040, 0.1857, 0.2355],\n",
       "            [0.1892, 0.1913, 0.1949, 0.1654, 0.2591],\n",
       "            [0.1678, 0.2405, 0.2269, 0.1539, 0.2109],\n",
       "            [0.2025, 0.2242, 0.2116, 0.1385, 0.2232]],\n",
       "  \n",
       "           [[0.2329, 0.1462, 0.1283, 0.3238, 0.1687],\n",
       "            [0.2043, 0.1700, 0.1572, 0.3385, 0.1301],\n",
       "            [0.1680, 0.1821, 0.1746, 0.3416, 0.1337],\n",
       "            [0.1632, 0.1754, 0.1780, 0.3523, 0.1311],\n",
       "            [0.1625, 0.1691, 0.1628, 0.3342, 0.1715]],\n",
       "  \n",
       "           [[0.1478, 0.2466, 0.2524, 0.1539, 0.1993],\n",
       "            [0.1602, 0.2590, 0.2339, 0.1350, 0.2118],\n",
       "            [0.1675, 0.2629, 0.2365, 0.1218, 0.2113],\n",
       "            [0.1301, 0.3145, 0.2946, 0.0882, 0.1726],\n",
       "            [0.1341, 0.3074, 0.2767, 0.1219, 0.1599]],\n",
       "  \n",
       "           [[0.2259, 0.1753, 0.1648, 0.2009, 0.2331],\n",
       "            [0.2487, 0.1447, 0.1301, 0.2140, 0.2624],\n",
       "            [0.2490, 0.1571, 0.1425, 0.2039, 0.2475],\n",
       "            [0.2397, 0.1849, 0.1695, 0.1903, 0.2156],\n",
       "            [0.2325, 0.1592, 0.1667, 0.2018, 0.2398]]]])],\n",
       " ['<bos>', 'C', 'C', 'O', '<eos>'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_full_attention(\"CCO\", model, DEFAULTS, TOKENIZER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = '/home/mi52qep/proj/molformer_refactor/src/molformer/data/bace/test.csv'\n",
    "BACE_TEST = pd.read_csv(DATA_PATH).smiles\n",
    "CHECKPOINT_PATHS = glob.glob(\"/home/mi52qep/proj/molformer/data/Pretrained MoLFormer/checkpoints/*.ckpt\")\n",
    "VOCAB_PATH = '/home/mi52qep/proj/molformer_refactor/src/molformer/data/bert_vocab.txt'\n",
    "\n",
    "\n",
    "CONFIG_FILE_PATH = '/home/mi52qep/proj/molformer/data/Pretrained MoLFormer/hparams.yaml'\n",
    "CONFIG = read_config(CONFIG_FILE_PATH)\n",
    "CONFIG.seed_path = CHECKPOINT_PATH\n",
    "CONFIG.canonical = False\n",
    "CONFIG.mask = False\n",
    "\n",
    "TEST_SMILES = BACE_TEST[0]\n",
    "\n",
    "for CHECKPOINT_PATH in CHECKPOINT_PATHS:\n",
    "    model = get_full_model(CHECKPOINT_PATH, VOCAB_PATH, CONFIG_FILE_PATH)\n",
    "    attention = agg_attention(TEST_SMILES, 12, CONFIG, VOCAB_PATH)\n",
    "    distance_matrix = distance_matrix(TEST_SMILES)\n",
    "    r2 = r2(attention, distance_matrix)\n",
    "    %%capture\n",
    "    print(r2)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "molformer",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}