qat_scratch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "950c2ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch_tensorrt\n",
    "from pytorch_quantization import quant_modules\n",
    "from models.experimental import attempt_load\n",
    "from utils.utils_qat import *\n",
    "from utils.datasets import create_dataloader\n",
    "import yaml\n",
    "from utils.general import colorstr\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b04e378d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fusing layers... \n",
      "IDetect.fuse\n"
     ]
    }
   ],
   "source": [
    "w = './runs/train/yolov7-tiny/weights/best.pt'\n",
    "# Load model\n",
    "quant_modules.initialize() # QAT\n",
    "torch_model = attempt_load(w).eval().cuda()\n",
    "torch_model.model[-1].training = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "48f07a61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[[[[-4.87547e-01,  9.47899e-02,  5.22560e-01,  ..., -6.50625e+00, -6.23370e+00, -4.29142e+00],\n",
       "            [-2.87736e-01, -1.69391e-01,  7.51416e-01,  ..., -6.16169e+00, -5.12955e+00, -3.47082e+00],\n",
       "            [-5.59410e-01, -1.65552e-02,  9.41929e-01,  ..., -6.76621e+00, -5.27977e+00, -4.36671e+00],\n",
       "            ...,\n",
       "            [-8.52922e-01, -1.31474e-02,  7.39594e-01,  ..., -7.51798e+00, -5.87583e+00, -4.87661e+00],\n",
       "            [-3.33085e-01,  9.42448e-03,  7.67701e-01,  ..., -8.03420e+00, -6.72067e+00, -5.64604e+00],\n",
       "            [-8.26194e-01,  8.08707e-02,  1.12668e+00,  ..., -5.60131e+00, -5.30826e+00, -3.10908e+00]],\n",
       " \n",
       "           [[-2.08714e-01, -1.34601e+00,  6.99600e-01,  ..., -4.64982e+00, -4.97741e+00, -2.42952e+00],\n",
       "            [-4.18270e-01, -6.20992e-01,  1.00484e+00,  ..., -5.67229e+00, -4.35944e+00, -3.12736e+00],\n",
       "            [ 1.66784e-02, -6.98056e-01,  9.36961e-01,  ..., -6.31265e+00, -5.18270e+00, -4.00719e+00],\n",
       "            ...,\n",
       "            [-1.08629e+00, -1.26781e-01,  7.61561e-01,  ..., -7.91392e+00, -6.29250e+00, -5.93685e+00],\n",
       "            [-3.89920e-01, -1.01824e+00,  9.75311e-01,  ..., -7.07589e+00, -5.45752e+00, -4.17389e+00],\n",
       "            [-9.42871e-04, -7.91197e-01,  7.01803e-01,  ..., -6.15737e+00, -5.63771e+00, -3.28445e+00]],\n",
       " \n",
       "           [[-4.14367e-01, -6.06221e-01,  5.22869e-01,  ..., -5.01067e+00, -4.24143e+00, -2.86208e+00],\n",
       "            [-5.17618e-03, -1.45199e-01,  8.91067e-01,  ..., -5.26744e+00, -3.89035e+00, -3.25030e+00],\n",
       "            [-1.30138e-01, -3.15366e-01,  8.70080e-01,  ..., -6.26594e+00, -4.48193e+00, -3.87450e+00],\n",
       "            ...,\n",
       "            [-4.44998e-01, -4.80577e-01,  6.80634e-01,  ..., -7.29893e+00, -5.58920e+00, -5.01736e+00],\n",
       "            [-3.73937e-01, -3.08182e-01,  8.17241e-01,  ..., -6.59318e+00, -5.17536e+00, -3.60148e+00],\n",
       "            [-1.49220e-01, -7.14043e-01,  5.57279e-01,  ..., -5.68227e+00, -4.91781e+00, -2.35666e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.94971e-01,  1.28529e-01,  5.07074e-01,  ..., -5.09467e+00, -3.67892e+00, -3.19462e+00],\n",
       "            [-2.68355e-01,  1.57400e-01,  9.50977e-01,  ..., -5.43986e+00, -3.81346e+00, -3.70236e+00],\n",
       "            [-4.14666e-02,  4.47814e-01,  9.39366e-01,  ..., -5.96580e+00, -3.85279e+00, -4.13694e+00],\n",
       "            ...,\n",
       "            [-1.07452e+00, -2.18644e-02,  4.55240e-01,  ..., -6.91368e+00, -5.13823e+00, -4.29441e+00],\n",
       "            [-8.05978e-01,  2.79831e-01,  5.62612e-01,  ..., -6.14632e+00, -4.68545e+00, -3.69438e+00],\n",
       "            [ 3.81541e-02, -4.45101e-01,  2.48568e-01,  ..., -5.42997e+00, -4.42222e+00, -3.16599e+00]],\n",
       " \n",
       "           [[-2.64487e-01,  8.11415e-01,  5.03177e-01,  ..., -6.06476e+00, -5.30296e+00, -3.55035e+00],\n",
       "            [-6.34020e-01,  1.26575e+00,  9.08688e-01,  ..., -6.47612e+00, -4.97852e+00, -4.37035e+00],\n",
       "            [-3.33445e-01,  1.40988e+00,  8.02741e-01,  ..., -6.79821e+00, -5.09136e+00, -4.69395e+00],\n",
       "            ...,\n",
       "            [-6.24214e-01,  4.19062e-01,  5.24241e-01,  ..., -6.66083e+00, -5.74247e+00, -4.12677e+00],\n",
       "            [-7.27590e-01,  6.69137e-01,  4.07194e-01,  ..., -6.44302e+00, -5.63481e+00, -3.75959e+00],\n",
       "            [ 3.21439e-02,  8.12609e-01,  4.12659e-01,  ..., -5.36296e+00, -4.98514e+00, -3.25482e+00]],\n",
       " \n",
       "           [[-2.31469e-01, -2.57948e-01,  6.51369e-01,  ..., -7.26445e+00, -6.55701e+00, -4.44299e+00],\n",
       "            [-1.03936e+00,  2.33757e-02,  1.16433e+00,  ..., -6.57010e+00, -5.77444e+00, -2.86514e+00],\n",
       "            [-5.24687e-01, -7.37154e-02,  8.11519e-01,  ..., -6.86183e+00, -6.18328e+00, -3.27283e+00],\n",
       "            ...,\n",
       "            [-1.16036e+00,  4.09528e-02,  5.30606e-01,  ..., -6.29941e+00, -6.00583e+00, -2.67043e+00],\n",
       "            [-8.65109e-01, -1.95316e-01,  4.96989e-01,  ..., -5.92568e+00, -5.88166e+00, -1.79793e+00],\n",
       "            [-6.76638e-01, -3.68356e-01,  4.90253e-01,  ..., -5.65889e+00, -5.37580e+00, -2.03957e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 3.15849e-02,  2.28550e-01,  3.46123e-01,  ..., -6.60015e+00, -6.58985e+00, -3.41172e+00],\n",
       "            [ 7.98145e-01, -3.09303e-02,  6.26568e-01,  ..., -6.42159e+00, -5.41520e+00, -2.64704e+00],\n",
       "            [ 5.87487e-01,  2.46096e-01,  7.25815e-01,  ..., -7.32118e+00, -5.57777e+00, -3.31774e+00],\n",
       "            ...,\n",
       "            [ 4.57211e-01,  9.63801e-02,  5.11913e-01,  ..., -8.40944e+00, -6.10431e+00, -3.85385e+00],\n",
       "            [ 4.97155e-01,  2.39678e-01,  4.38248e-01,  ..., -8.85266e+00, -7.09493e+00, -4.93373e+00],\n",
       "            [-5.09883e-02,  2.67987e-01,  7.14785e-01,  ..., -5.18850e+00, -5.50191e+00, -2.56686e+00]],\n",
       " \n",
       "           [[-9.97706e-02, -9.72307e-01,  6.74504e-01,  ..., -4.40776e+00, -5.17882e+00, -1.90637e+00],\n",
       "            [ 6.89782e-01, -5.56975e-01,  1.30114e+00,  ..., -5.80840e+00, -4.83887e+00, -3.06389e+00],\n",
       "            [ 6.63059e-01, -2.09686e-01,  1.00001e+00,  ..., -6.74464e+00, -5.21759e+00, -3.67518e+00],\n",
       "            ...,\n",
       "            [-2.68900e-01,  2.71106e-02,  6.59916e-01,  ..., -9.17111e+00, -6.36998e+00, -5.88036e+00],\n",
       "            [ 8.75210e-01, -9.10337e-01,  9.88143e-01,  ..., -7.40403e+00, -5.85184e+00, -3.73080e+00],\n",
       "            [ 1.48203e-01, -3.83418e-01,  4.23807e-01,  ..., -6.06164e+00, -5.40158e+00, -3.06784e+00]],\n",
       " \n",
       "           [[-1.37933e-01, -3.14333e-01,  5.39705e-01,  ..., -4.98543e+00, -4.11740e+00, -2.54125e+00],\n",
       "            [ 7.18334e-01, -7.52797e-02,  1.12358e+00,  ..., -5.19707e+00, -3.97505e+00, -3.04478e+00],\n",
       "            [ 8.29768e-01, -1.53299e-01,  9.16135e-01,  ..., -6.70338e+00, -4.67571e+00, -3.71698e+00],\n",
       "            ...,\n",
       "            [ 2.41153e-01, -1.51111e-01,  5.98454e-01,  ..., -8.18722e+00, -5.74897e+00, -4.89638e+00],\n",
       "            [ 5.16505e-01, -2.94282e-03,  9.88372e-01,  ..., -7.07471e+00, -5.27063e+00, -3.38013e+00],\n",
       "            [ 3.05060e-01, -3.41958e-01,  3.28326e-01,  ..., -5.62924e+00, -4.81111e+00, -2.16957e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.76056e-01,  7.25818e-01,  5.78552e-01,  ..., -5.09775e+00, -3.23365e+00, -3.13971e+00],\n",
       "            [ 5.68397e-01,  6.26875e-01,  1.22123e+00,  ..., -5.67964e+00, -3.96818e+00, -3.51433e+00],\n",
       "            [ 1.08453e+00,  6.45110e-01,  1.12883e+00,  ..., -6.31458e+00, -3.83670e+00, -3.83725e+00],\n",
       "            ...,\n",
       "            [ 4.47320e-02,  2.63744e-01,  4.18985e-01,  ..., -7.91338e+00, -5.27141e+00, -4.18778e+00],\n",
       "            [ 5.50769e-01,  6.68986e-01,  5.05343e-01,  ..., -7.11973e+00, -4.86756e+00, -3.73740e+00],\n",
       "            [ 4.33117e-01, -2.00450e-02,  5.76506e-02,  ..., -5.56409e+00, -4.26651e+00, -3.08853e+00]],\n",
       " \n",
       "           [[ 6.97941e-02,  1.42860e+00,  5.09332e-01,  ..., -6.66475e+00, -5.26134e+00, -2.87214e+00],\n",
       "            [ 9.50227e-01,  1.57617e+00,  9.80701e-01,  ..., -7.18167e+00, -5.11077e+00, -3.70764e+00],\n",
       "            [ 6.60737e-01,  1.82962e+00,  7.08275e-01,  ..., -7.75289e+00, -4.98907e+00, -3.96828e+00],\n",
       "            ...,\n",
       "            [-2.38761e-01,  7.89576e-01,  4.34184e-01,  ..., -7.57988e+00, -5.72025e+00, -3.90694e+00],\n",
       "            [ 3.42054e-01,  1.04097e+00,  3.83656e-01,  ..., -7.45516e+00, -5.88741e+00, -3.63914e+00],\n",
       "            [ 3.51212e-01,  1.13755e+00,  2.04610e-01,  ..., -5.51940e+00, -5.01503e+00, -3.01027e+00]],\n",
       " \n",
       "           [[ 1.41553e-01, -3.22121e-01,  4.03034e-01,  ..., -7.84405e+00, -6.58219e+00, -4.06551e+00],\n",
       "            [ 8.11673e-01, -1.70145e-01,  9.59076e-01,  ..., -7.34860e+00, -5.79930e+00, -2.74612e+00],\n",
       "            [ 5.01517e-01, -1.39298e-01,  5.73288e-01,  ..., -7.61431e+00, -6.30421e+00, -3.19174e+00],\n",
       "            ...,\n",
       "            [ 8.57110e-02, -8.38753e-03,  3.65468e-01,  ..., -6.82223e+00, -6.07952e+00, -2.66817e+00],\n",
       "            [ 3.48529e-01, -1.39417e-01,  4.36313e-01,  ..., -6.81344e+00, -6.08143e+00, -1.78918e+00],\n",
       "            [ 2.00045e-01, -3.38761e-01,  2.91028e-01,  ..., -6.06565e+00, -5.54764e+00, -1.91761e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.32191e-01,  3.84730e-01, -1.92219e-01,  ..., -5.19991e+00, -6.19433e+00, -2.18451e+00],\n",
       "            [ 1.99610e-01,  4.29013e-02,  3.51371e-01,  ..., -4.83985e+00, -4.62759e+00, -1.85931e+00],\n",
       "            [-6.33848e-02,  2.48728e-01,  2.21740e-01,  ..., -5.92031e+00, -4.70049e+00, -2.95003e+00],\n",
       "            ...,\n",
       "            [-2.43293e-01,  7.83563e-02,  1.26316e-01,  ..., -6.85647e+00, -5.07950e+00, -2.73201e+00],\n",
       "            [-9.91466e-02,  2.50978e-01,  1.77051e-02,  ..., -7.19382e+00, -6.21510e+00, -3.55181e+00],\n",
       "            [-4.87552e-01,  3.68024e-01,  2.29515e-01,  ..., -3.78722e+00, -5.14449e+00, -1.97985e+00]],\n",
       " \n",
       "           [[-4.62441e-02, -4.40942e-01, -3.57341e-02,  ..., -3.76818e+00, -5.00800e+00, -2.38453e+00],\n",
       "            [-1.59180e-01, -2.64265e-01,  7.75303e-01,  ..., -5.10047e+00, -4.02997e+00, -3.77329e+00],\n",
       "            [ 1.17001e-01,  2.35244e-01,  5.25368e-01,  ..., -5.91765e+00, -4.40940e+00, -4.28323e+00],\n",
       "            ...,\n",
       "            [-6.52011e-01,  1.56518e-01,  1.33104e-01,  ..., -8.13151e+00, -5.45183e+00, -4.92595e+00],\n",
       "            [-5.46222e-02, -6.67767e-01,  4.49806e-01,  ..., -6.37625e+00, -4.94734e+00, -2.87220e+00],\n",
       "            [-1.09597e-01, -8.53180e-02,  5.02387e-02,  ..., -5.00821e+00, -5.11328e+00, -2.39618e+00]],\n",
       " \n",
       "           [[-8.71003e-02,  3.34449e-02, -6.59637e-02,  ..., -4.05661e+00, -3.98292e+00, -2.86342e+00],\n",
       "            [ 2.30050e-01,  1.76727e-01,  6.03876e-01,  ..., -4.26069e+00, -3.45791e+00, -3.47130e+00],\n",
       "            [ 3.06481e-02,  1.44447e-01,  4.08828e-01,  ..., -5.93517e+00, -3.93198e+00, -4.32104e+00],\n",
       "            ...,\n",
       "            [-2.23589e-01,  1.78737e-01,  8.13293e-02,  ..., -7.43399e+00, -4.92224e+00, -4.75224e+00],\n",
       "            [-2.37804e-01,  3.81797e-01,  4.62065e-01,  ..., -6.25903e+00, -4.47447e+00, -3.46829e+00],\n",
       "            [-1.11637e-01,  1.41449e-04, -8.56984e-03,  ..., -4.79554e+00, -4.30185e+00, -1.84170e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.90861e-02,  1.19377e+00,  2.90935e-03,  ..., -3.97029e+00, -3.33947e+00, -3.57183e+00],\n",
       "            [ 5.41559e-02,  1.17661e+00,  6.60975e-01,  ..., -4.85459e+00, -3.39240e+00, -4.14737e+00],\n",
       "            [ 3.66612e-01,  1.00411e+00,  6.48777e-01,  ..., -5.64053e+00, -3.10959e+00, -4.19861e+00],\n",
       "            ...,\n",
       "            [-5.44627e-01,  7.45990e-01, -5.96765e-02,  ..., -7.43756e+00, -4.25536e+00, -4.66371e+00],\n",
       "            [-9.07632e-02,  9.66968e-01,  1.17213e-01,  ..., -6.46287e+00, -3.87151e+00, -4.39357e+00],\n",
       "            [ 8.13054e-02,  3.87867e-01, -1.93750e-01,  ..., -4.96531e+00, -3.87860e+00, -3.29024e+00]],\n",
       " \n",
       "           [[ 9.83985e-02,  1.97879e+00, -1.61114e-01,  ..., -5.75974e+00, -4.90095e+00, -3.08028e+00],\n",
       "            [ 4.00071e-02,  1.84634e+00,  4.66661e-01,  ..., -6.31342e+00, -4.29636e+00, -4.37891e+00],\n",
       "            [-3.89413e-02,  2.20738e+00,  1.16960e-01,  ..., -7.02627e+00, -4.36548e+00, -5.26443e+00],\n",
       "            ...,\n",
       "            [-5.01535e-01,  1.17850e+00, -9.57908e-02,  ..., -6.95734e+00, -4.97889e+00, -4.12904e+00],\n",
       "            [-2.36574e-01,  1.31699e+00, -1.60922e-01,  ..., -6.75537e+00, -5.09192e+00, -4.28928e+00],\n",
       "            [ 1.69769e-02,  1.39285e+00, -2.22088e-01,  ..., -4.90206e+00, -4.62363e+00, -3.13600e+00]],\n",
       " \n",
       "           [[ 5.55691e-02, -2.36844e-01, -2.22252e-01,  ..., -7.09046e+00, -6.49009e+00, -3.78842e+00],\n",
       "            [-1.59989e-01, -9.94457e-02,  3.79156e-01,  ..., -6.56837e+00, -5.15566e+00, -2.75006e+00],\n",
       "            [-7.64427e-02, -5.52362e-02, -3.25950e-02,  ..., -6.96484e+00, -5.91400e+00, -3.61679e+00],\n",
       "            ...,\n",
       "            [-4.26746e-01,  5.46663e-02, -1.61771e-01,  ..., -5.92453e+00, -5.60054e+00, -2.66462e+00],\n",
       "            [-3.59739e-01,  7.96747e-03, -5.02833e-02,  ..., -6.23178e+00, -5.54439e+00, -2.09229e+00],\n",
       "            [-3.41345e-01, -1.58885e-01, -1.59096e-01,  ..., -5.44165e+00, -5.20717e+00, -2.09562e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-5.34644e-01, -2.11086e-01, -2.01626e-02,  ..., -7.82099e+00, -2.49750e+00, -6.42748e+00],\n",
       "            [ 1.42874e-01, -2.77607e-01,  9.56378e-01,  ..., -7.81695e+00, -2.44013e+00, -5.85250e+00],\n",
       "            [-2.86203e-01, -4.66548e-01,  9.10675e-01,  ..., -7.76322e+00, -2.97811e+00, -5.65296e+00],\n",
       "            ...,\n",
       "            [ 3.68874e-01, -9.35728e-02,  1.34625e+00,  ..., -5.84471e+00, -3.61456e+00, -5.03081e+00],\n",
       "            [ 1.54349e-01, -3.24116e-02,  9.09287e-01,  ..., -6.48826e+00, -4.25515e+00, -5.64304e+00],\n",
       "            [ 1.04212e-01, -1.31842e-01,  7.13947e-01,  ..., -5.52953e+00, -2.89049e+00, -5.82240e+00]],\n",
       " \n",
       "           [[-5.98689e-01,  2.14755e-02, -9.57832e-02,  ..., -7.97864e+00, -3.23367e+00, -5.88555e+00],\n",
       "            [ 6.29052e-01, -2.91316e-01,  5.06456e-01,  ..., -9.68456e+00, -4.64034e+00, -6.17084e+00],\n",
       "            [-3.31653e-01,  3.08612e-01,  5.68547e-01,  ..., -9.90234e+00, -4.04214e+00, -6.37091e+00],\n",
       "            ...,\n",
       "            [-8.43346e-02,  4.71697e-01,  6.66470e-01,  ..., -7.96878e+00, -4.89778e+00, -5.51477e+00],\n",
       "            [-3.60978e-01,  1.17822e+00,  5.38258e-01,  ..., -7.85644e+00, -5.40036e+00, -6.15396e+00],\n",
       "            [ 2.34915e-01, -6.51149e-01,  2.54822e-01,  ..., -6.24797e+00, -3.71015e+00, -5.53734e+00]],\n",
       " \n",
       "           [[-3.85954e-01,  8.37675e-01,  1.01079e-01,  ..., -9.24176e+00, -5.42115e+00, -6.93347e+00],\n",
       "            [-7.47709e-01, -1.99777e-01,  4.66645e-01,  ..., -1.02493e+01, -4.84742e+00, -6.17880e+00],\n",
       "            [ 3.74259e-01,  1.94674e-01,  3.79761e-01,  ..., -1.11790e+01, -4.89673e+00, -6.55665e+00],\n",
       "            ...,\n",
       "            [ 2.57704e-01, -3.11711e-01,  3.17036e-01,  ..., -8.93264e+00, -5.72079e+00, -6.15489e+00],\n",
       "            [-9.66079e-02, -4.40850e-01,  2.85533e-01,  ..., -9.99869e+00, -5.79011e+00, -6.97508e+00],\n",
       "            [-6.84820e-01, -1.14070e-01,  4.35575e-01,  ..., -7.80051e+00, -4.08738e+00, -6.01042e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.53967e-01,  1.88052e-01, -1.23030e-01,  ..., -9.03535e+00, -4.23133e+00, -6.26092e+00],\n",
       "            [ 3.54678e-01,  5.20671e-01,  3.16505e-01,  ..., -1.06015e+01, -6.02396e+00, -6.65873e+00],\n",
       "            [-1.06829e+00,  1.98856e-01,  5.94994e-01,  ..., -9.14564e+00, -4.63469e+00, -5.87122e+00],\n",
       "            ...,\n",
       "            [-3.20582e-01, -6.96720e-01,  3.28711e-01,  ..., -8.96921e+00, -5.89011e+00, -5.91327e+00],\n",
       "            [-3.67608e-01,  4.71955e-01,  1.80950e-01,  ..., -1.04861e+01, -6.42393e+00, -6.25397e+00],\n",
       "            [ 3.04558e-01, -1.39310e-01,  2.07510e-01,  ..., -7.38946e+00, -3.09614e+00, -5.63271e+00]],\n",
       " \n",
       "           [[-6.03119e-01, -9.63672e-02, -1.04045e-01,  ..., -8.04713e+00, -4.58191e+00, -6.08050e+00],\n",
       "            [ 2.59290e-01, -9.73837e-01,  2.25739e-01,  ..., -1.07871e+01, -5.90256e+00, -7.13039e+00],\n",
       "            [-3.28928e-01,  1.13634e+00,  5.37875e-01,  ..., -9.05449e+00, -4.34472e+00, -5.74742e+00],\n",
       "            ...,\n",
       "            [ 5.64289e-01,  1.98350e-01,  3.47060e-01,  ..., -9.01040e+00, -4.89771e+00, -5.81461e+00],\n",
       "            [-6.64502e-01, -7.74392e-02,  1.19169e-01,  ..., -9.52879e+00, -5.48731e+00, -6.12110e+00],\n",
       "            [ 3.87243e-02,  1.22087e+00,  4.11300e-01,  ..., -6.64341e+00, -2.98357e+00, -5.37229e+00]],\n",
       " \n",
       "           [[-3.68444e-01,  2.45559e-01,  2.05618e-01,  ..., -7.28658e+00, -5.36106e+00, -7.13406e+00],\n",
       "            [-2.28250e-01,  2.67206e-01,  4.23454e-01,  ..., -8.43273e+00, -5.07172e+00, -6.90686e+00],\n",
       "            [-8.73098e-01,  2.01464e-01,  4.56206e-01,  ..., -8.45828e+00, -3.64197e+00, -6.49234e+00],\n",
       "            ...,\n",
       "            [ 3.53060e-01,  3.73122e-01,  2.77468e-01,  ..., -7.86769e+00, -4.67534e+00, -6.31371e+00],\n",
       "            [-7.49054e-02,  1.28421e-01,  1.24344e-01,  ..., -7.36432e+00, -4.73637e+00, -6.40824e+00],\n",
       "            [ 8.29728e-03,  9.05656e-02,  2.18522e-01,  ..., -5.74112e+00, -3.55244e+00, -6.04109e+00]]],\n",
       " \n",
       " \n",
       "          [[[-4.06517e-01, -3.13839e-01, -4.32238e-01,  ..., -5.68273e+00, -1.77525e+00, -4.28525e+00],\n",
       "            [ 2.32293e-01, -3.82456e-01,  3.46532e-01,  ..., -5.17990e+00, -1.99043e+00, -5.60844e+00],\n",
       "            [-2.72103e-01, -5.71252e-01,  3.48941e-01,  ..., -5.26630e+00, -1.84904e+00, -5.00860e+00],\n",
       "            ...,\n",
       "            [ 3.21151e-01, -2.30735e-01,  6.15465e-01,  ..., -4.78941e+00, -3.23944e+00, -4.22455e+00],\n",
       "            [ 1.35220e-01, -1.48121e-01,  3.44990e-01,  ..., -5.05759e+00, -3.35567e+00, -4.01805e+00],\n",
       "            [-9.62100e-03, -2.05407e-01,  1.27524e-02,  ..., -5.52845e+00, -3.02920e+00, -3.12888e+00]],\n",
       " \n",
       "           [[-5.50994e-01, -2.07030e-01, -5.20789e-01,  ..., -5.36434e+00, -2.55363e+00, -4.77837e+00],\n",
       "            [ 6.79680e-01, -3.73727e-01,  2.80169e-02,  ..., -5.12977e+00, -2.84838e+00, -6.32173e+00],\n",
       "            [-3.39791e-01,  2.20178e-01,  3.41156e-02,  ..., -5.06108e+00, -2.18273e+00, -6.41758e+00],\n",
       "            ...,\n",
       "            [-1.12713e-01,  3.39080e-01,  2.91680e-01,  ..., -4.67470e+00, -3.77412e+00, -5.01133e+00],\n",
       "            [-4.51763e-01,  1.04552e+00,  1.39919e-01,  ..., -4.95790e+00, -4.18028e+00, -4.61391e+00],\n",
       "            [ 1.30261e-01, -7.65899e-01, -1.46250e-01,  ..., -5.30267e+00, -3.19884e+00, -4.02355e+00]],\n",
       " \n",
       "           [[-3.84929e-01,  7.28741e-01, -4.45305e-01,  ..., -5.78910e+00, -3.51964e+00, -4.07149e+00],\n",
       "            [-8.21534e-01, -3.43537e-01, -6.48426e-03,  ..., -5.28001e+00, -2.71841e+00, -6.01262e+00],\n",
       "            [ 3.66150e-01,  1.21018e-01, -1.06943e-01,  ..., -5.24157e+00, -2.66553e+00, -7.21302e+00],\n",
       "            ...,\n",
       "            [ 3.11385e-01, -3.41815e-01, -1.88695e-02,  ..., -5.13963e+00, -3.76725e+00, -4.57831e+00],\n",
       "            [-5.35542e-02, -4.73546e-01, -1.37965e-01,  ..., -5.69140e+00, -3.94770e+00, -3.54279e+00],\n",
       "            [-7.94450e-01, -1.50748e-01, -5.05972e-02,  ..., -5.44475e+00, -3.25262e+00, -3.60718e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.70728e-01,  1.19857e-01, -5.07992e-01,  ..., -5.36931e+00, -3.07238e+00, -5.90686e+00],\n",
       "            [ 3.13654e-01,  4.44815e-01, -2.13364e-01,  ..., -5.37257e+00, -3.83852e+00, -6.50270e+00],\n",
       "            [-1.23727e+00,  4.44761e-02,  3.07791e-02,  ..., -4.99012e+00, -3.33906e+00, -6.62603e+00],\n",
       "            ...,\n",
       "            [-3.88726e-01, -6.81882e-01, -5.67368e-03,  ..., -5.21564e+00, -3.74422e+00, -5.28564e+00],\n",
       "            [-4.10594e-01,  4.17406e-01, -1.79837e-01,  ..., -5.26571e+00, -3.86466e+00, -5.89214e+00],\n",
       "            [ 2.47219e-01, -2.16870e-01, -2.61808e-01,  ..., -5.17811e+00, -2.65242e+00, -6.01686e+00]],\n",
       " \n",
       "           [[-5.02138e-01, -1.63649e-01, -4.54308e-01,  ..., -5.48369e+00, -3.73103e+00, -4.08149e+00],\n",
       "            [ 2.66751e-01, -1.00793e+00, -2.47442e-01,  ..., -6.02199e+00, -4.07716e+00, -5.40449e+00],\n",
       "            [-4.26859e-01,  8.34247e-01,  4.43697e-02,  ..., -4.80205e+00, -3.39179e+00, -5.50884e+00],\n",
       "            ...,\n",
       "            [ 5.64258e-01,  1.37289e-01, -1.72616e-02,  ..., -4.77228e+00, -3.12723e+00, -5.06410e+00],\n",
       "            [-5.92736e-01, -1.03649e-01, -2.51549e-01,  ..., -5.33272e+00, -3.71323e+00, -4.45815e+00],\n",
       "            [-3.00860e-02,  1.06360e+00, -9.15811e-02,  ..., -4.79556e+00, -3.29877e+00, -5.69029e+00]],\n",
       " \n",
       "           [[-1.38034e-01,  2.05968e-01, -1.84267e-01,  ..., -6.12671e+00, -4.33077e+00, -1.75113e+00],\n",
       "            [-2.05111e-01,  1.83119e-01,  1.11804e-01,  ..., -5.73394e+00, -3.94210e+00, -4.18133e+00],\n",
       "            [-1.03115e+00,  1.31382e-01,  5.55934e-02,  ..., -5.44817e+00, -2.74345e+00, -4.35958e+00],\n",
       "            ...,\n",
       "            [ 3.85629e-01,  3.63764e-01, -2.63131e-02,  ..., -5.45473e+00, -3.52178e+00, -3.46353e+00],\n",
       "            [ 9.22192e-02,  1.24744e-01, -1.25533e-01,  ..., -5.67737e+00, -4.07566e+00, -2.21402e+00],\n",
       "            [-9.37037e-02, -7.55376e-04, -1.02389e-01,  ..., -5.61605e+00, -3.50843e+00, -2.99084e+00]]],\n",
       " \n",
       " \n",
       "          [[[-3.93481e-01, -1.64835e-01, -2.65352e-01,  ..., -6.24092e+00, -1.37638e+00, -4.65687e+00],\n",
       "            [ 2.36445e-01, -2.17818e-01,  6.65350e-01,  ..., -6.50355e+00, -1.58506e+00, -5.21347e+00],\n",
       "            [-2.64106e-01, -3.86140e-01,  6.97016e-01,  ..., -6.37442e+00, -1.73013e+00, -6.06838e+00],\n",
       "            ...,\n",
       "            [ 3.23036e-01, -2.81005e-02,  1.02931e+00,  ..., -4.98665e+00, -2.67558e+00, -4.78585e+00],\n",
       "            [ 1.36195e-01,  3.15027e-02,  6.86678e-01,  ..., -5.27352e+00, -3.31995e+00, -4.32447e+00],\n",
       "            [-1.41374e-02, -6.82866e-02,  2.60394e-01,  ..., -4.67739e+00, -2.04981e+00, -3.33594e+00]],\n",
       " \n",
       "           [[-5.31133e-01, -6.89732e-02, -3.79144e-01,  ..., -6.24209e+00, -2.04830e+00, -4.62131e+00],\n",
       "            [ 7.02004e-01, -3.40965e-01,  2.80657e-01,  ..., -7.65522e+00, -3.40918e+00, -6.82172e+00],\n",
       "            [-3.11834e-01,  2.66643e-01,  2.85398e-01,  ..., -8.31331e+00, -2.58867e+00, -8.60124e+00],\n",
       "            ...,\n",
       "            [-9.23127e-02,  3.59382e-01,  6.58524e-01,  ..., -6.21727e+00, -3.77158e+00, -5.79024e+00],\n",
       "            [-4.31768e-01,  1.09399e+00,  4.12277e-01,  ..., -6.15885e+00, -4.46057e+00, -4.96148e+00],\n",
       "            [ 1.27948e-01, -6.79140e-01,  1.24043e-01,  ..., -5.27131e+00, -2.77939e+00, -3.89982e+00]],\n",
       " \n",
       "           [[-3.75162e-01,  8.02944e-01, -3.22589e-01,  ..., -6.99978e+00, -3.84780e+00, -6.32713e+00],\n",
       "            [-7.99745e-01, -3.01548e-01,  2.11738e-01,  ..., -8.15942e+00, -2.94807e+00, -8.73102e+00],\n",
       "            [ 3.89190e-01,  1.43063e-01,  1.56118e-01,  ..., -8.94957e+00, -2.92532e+00, -9.84021e+00],\n",
       "            ...,\n",
       "            [ 3.36719e-01, -3.41774e-01,  2.19717e-01,  ..., -7.17027e+00, -4.64166e+00, -5.34516e+00],\n",
       "            [-4.04283e-02, -4.84423e-01,  7.24958e-02,  ..., -7.53929e+00, -4.57006e+00, -4.88177e+00],\n",
       "            [-7.71968e-01, -8.76201e-02,  2.07294e-01,  ..., -5.97681e+00, -3.06944e+00, -3.92303e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.56606e-01,  1.48564e-01, -3.71658e-01,  ..., -7.08749e+00, -2.97126e+00, -4.88778e+00],\n",
       "            [ 3.28686e-01,  4.48683e-01, -1.60536e-02,  ..., -7.97802e+00, -4.48025e+00, -6.42402e+00],\n",
       "            [-1.22843e+00,  8.45532e-02,  2.90987e-01,  ..., -7.46568e+00, -3.21178e+00, -8.14664e+00],\n",
       "            ...,\n",
       "            [-3.59539e-01, -6.81710e-01,  2.80953e-01,  ..., -7.61724e+00, -4.54187e+00, -7.37267e+00],\n",
       "            [-3.81460e-01,  4.56359e-01,  5.24283e-02,  ..., -8.28330e+00, -4.93384e+00, -6.63871e+00],\n",
       "            [ 2.58176e-01, -1.33873e-01, -1.88512e-02,  ..., -6.70662e+00, -2.29315e+00, -5.21662e+00]],\n",
       " \n",
       "           [[-4.78702e-01, -3.86433e-02, -3.30032e-01,  ..., -6.36707e+00, -3.81873e+00, -3.17556e+00],\n",
       "            [ 2.75605e-01, -1.00303e+00, -9.63290e-02,  ..., -8.13405e+00, -4.98759e+00, -4.68111e+00],\n",
       "            [-4.01605e-01,  9.37307e-01,  2.78397e-01,  ..., -7.13447e+00, -3.21912e+00, -6.31424e+00],\n",
       "            ...,\n",
       "            [ 5.86823e-01,  1.46871e-01,  2.39290e-01,  ..., -7.55669e+00, -3.65415e+00, -7.08223e+00],\n",
       "            [-5.59489e-01, -8.19582e-02, -1.31855e-02,  ..., -7.41801e+00, -4.03407e+00, -5.99809e+00],\n",
       "            [-2.62739e-02,  1.16120e+00,  1.83879e-01,  ..., -6.01141e+00, -2.43032e+00, -5.07096e+00]],\n",
       " \n",
       "           [[-1.31486e-01,  2.19938e-01, -2.80664e-02,  ..., -5.06480e+00, -4.41857e+00, -1.47263e+00],\n",
       "            [-1.90014e-01,  1.90943e-01,  3.87879e-01,  ..., -6.09805e+00, -4.07985e+00, -3.13213e+00],\n",
       "            [-1.02216e+00,  1.13174e-01,  3.36514e-01,  ..., -6.59328e+00, -2.49513e+00, -4.24345e+00],\n",
       "            ...,\n",
       "            [ 3.92900e-01,  3.41196e-01,  2.55738e-01,  ..., -6.29108e+00, -3.63153e+00, -3.24550e+00],\n",
       "            [ 1.09300e-01,  9.40300e-02,  1.62069e-01,  ..., -5.43772e+00, -3.66188e+00, -1.87217e+00],\n",
       "            [-8.76247e-02, -3.04074e-02,  1.73582e-01,  ..., -4.63921e+00, -2.82835e+00, -2.12553e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-3.93331e-01, -1.05222e-01, -5.09549e-01,  ..., -5.09263e+00, -2.47104e+00, -3.12491e+00],\n",
       "            [-4.36017e-01, -8.28347e-02, -2.37968e-01,  ..., -5.72009e+00, -4.77591e+00, -3.22002e+00],\n",
       "            [-5.58973e-01,  1.58730e-02, -5.06258e-01,  ..., -5.61712e+00, -6.18603e+00, -2.49051e+00],\n",
       "            ...,\n",
       "            [-4.39302e-01, -3.92893e-01, -3.87710e-02,  ..., -5.48526e+00, -5.74917e+00, -3.30261e+00],\n",
       "            [ 2.02982e-01, -6.98141e-01,  3.98241e-01,  ..., -5.69216e+00, -4.42809e+00, -4.64140e+00],\n",
       "            [ 7.33213e-02, -7.30040e-02, -1.24290e-01,  ..., -5.28591e+00, -3.06759e+00, -4.03527e+00]],\n",
       " \n",
       "           [[-6.06937e-01,  1.14335e-01, -8.62589e-01,  ..., -5.96034e+00, -5.82597e+00, -2.66099e+00],\n",
       "            [ 1.43754e-01, -6.31809e-01, -7.11080e-01,  ..., -6.13638e+00, -6.65993e+00, -3.67516e+00],\n",
       "            [ 4.14461e-01, -1.70360e-01, -6.52466e-01,  ..., -6.09841e+00, -7.22256e+00, -3.72888e+00],\n",
       "            ...,\n",
       "            [ 7.22500e-01,  1.09599e-01, -5.98348e-01,  ..., -6.43534e+00, -7.67048e+00, -3.09627e+00],\n",
       "            [ 6.49903e-01, -7.91262e-01, -5.43980e-01,  ..., -5.95212e+00, -6.60124e+00, -3.09896e+00],\n",
       "            [ 9.17420e-02, -1.90037e-01, -5.32984e-01,  ..., -5.41230e+00, -4.15834e+00, -2.67232e+00]],\n",
       " \n",
       "           [[-2.46508e-01, -2.39454e-01, -7.72591e-01,  ..., -5.80634e+00, -5.51145e+00, -2.65342e+00],\n",
       "            [-6.48271e-01,  7.46428e-02, -4.36441e-01,  ..., -6.56047e+00, -7.83517e+00, -4.28045e+00],\n",
       "            [-7.30683e-01, -9.16873e-02, -4.59250e-01,  ..., -6.24587e+00, -6.84750e+00, -3.84881e+00],\n",
       "            ...,\n",
       "            [ 2.59709e-01, -3.66455e-01, -2.33586e-01,  ..., -5.38156e+00, -4.86559e+00, -3.29966e+00],\n",
       "            [-1.47327e-01,  6.20925e-01, -6.87422e-01,  ..., -6.17832e+00, -7.19875e+00, -3.85422e+00],\n",
       "            [ 1.12484e-01, -3.03807e-01, -6.31626e-01,  ..., -5.51662e+00, -3.70880e+00, -3.29635e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.90662e-01, -7.56013e-02, -7.73741e-01,  ..., -5.59103e+00, -4.24718e+00, -2.80928e+00],\n",
       "            [ 8.22709e-01, -8.01424e-01, -6.01454e-01,  ..., -6.15319e+00, -6.31493e+00, -3.33469e+00],\n",
       "            [-1.38188e-01,  3.61699e-01, -3.78646e-01,  ..., -6.04368e+00, -6.68694e+00, -3.72436e+00],\n",
       "            ...,\n",
       "            [-3.87502e-01, -1.36233e-01, -5.83087e-01,  ..., -6.01715e+00, -7.46920e+00, -3.25095e+00],\n",
       "            [-6.70333e-01, -1.13094e-01, -3.23426e-01,  ..., -5.52235e+00, -6.41099e+00, -3.71119e+00],\n",
       "            [ 1.49265e-01, -5.59742e-01, -4.61881e-01,  ..., -5.25705e+00, -3.18051e+00, -3.62379e+00]],\n",
       " \n",
       "           [[-6.61439e-01,  6.91265e-02, -7.33034e-01,  ..., -5.32061e+00, -2.57251e+00, -2.91499e+00],\n",
       "            [-9.52094e-01,  7.14009e-02, -4.64346e-01,  ..., -6.00109e+00, -6.41508e+00, -3.61860e+00],\n",
       "            [ 4.35353e-02,  9.19674e-01,  2.35582e-02,  ..., -5.70810e+00, -5.96831e+00, -3.95499e+00],\n",
       "            ...,\n",
       "            [ 6.28351e-01,  7.17305e-02, -8.01041e-01,  ..., -6.43654e+00, -7.39855e+00, -3.87321e+00],\n",
       "            [ 8.01390e-01, -4.80151e-01, -3.57970e-01,  ..., -5.86190e+00, -6.20281e+00, -4.15007e+00],\n",
       "            [-1.11148e-01,  8.27017e-01, -3.31616e-01,  ..., -5.30484e+00, -4.09328e+00, -3.56802e+00]],\n",
       " \n",
       "           [[-2.06759e-01,  4.18588e-01, -2.41800e-01,  ..., -5.03750e+00, -2.88443e+00, -1.92594e+00],\n",
       "            [-1.22037e+00,  3.09087e-01, -7.56414e-02,  ..., -4.64065e+00, -3.93687e+00, -2.65132e+00],\n",
       "            [-8.41881e-02,  3.51969e-01,  1.16167e-01,  ..., -4.45952e+00, -4.29821e+00, -1.86746e+00],\n",
       "            ...,\n",
       "            [ 3.18738e-01,  7.70666e-01, -3.03295e-01,  ..., -5.16802e+00, -4.41680e+00, -1.89821e+00],\n",
       "            [ 2.10930e-01,  7.95963e-01, -1.15539e-01,  ..., -5.13598e+00, -5.09856e+00, -2.50403e+00],\n",
       "            [-1.27559e-01,  1.78974e-01, -1.65130e-01,  ..., -4.83877e+00, -4.01592e+00, -3.01496e+00]]],\n",
       " \n",
       " \n",
       "          [[[-4.00558e-01, -9.29807e-02, -4.43181e-01,  ..., -4.63613e+00, -3.86659e+00, -6.68300e+00],\n",
       "            [-4.26476e-01, -4.49553e-02, -1.92937e-01,  ..., -7.59619e+00, -6.35421e+00, -7.80962e+00],\n",
       "            [-5.54059e-01,  5.25383e-02, -5.38825e-01,  ..., -7.59989e+00, -7.04849e+00, -7.35448e+00],\n",
       "            ...,\n",
       "            [-4.29453e-01, -3.56062e-01,  2.68131e-02,  ..., -6.93618e+00, -6.78547e+00, -7.46180e+00],\n",
       "            [ 2.16549e-01, -6.80154e-01,  5.92998e-01,  ..., -5.96456e+00, -5.22801e+00, -7.29059e+00],\n",
       "            [ 8.90230e-02, -6.85379e-02,  1.57338e-03,  ..., -3.90567e+00, -4.11370e+00, -6.70960e+00]],\n",
       " \n",
       "           [[-6.14775e-01,  1.01046e-01, -8.43621e-01,  ..., -8.20973e+00, -7.13385e+00, -8.07922e+00],\n",
       "            [ 1.46915e-01, -6.37277e-01, -7.73883e-01,  ..., -9.83100e+00, -7.55096e+00, -8.28613e+00],\n",
       "            [ 4.18250e-01, -1.64201e-01, -7.25267e-01,  ..., -1.07829e+01, -8.22528e+00, -8.40421e+00],\n",
       "            ...,\n",
       "            [ 7.35726e-01,  9.82043e-02, -6.35903e-01,  ..., -1.09155e+01, -8.95914e+00, -8.72411e+00],\n",
       "            [ 6.57374e-01, -8.45952e-01, -5.66978e-01,  ..., -9.20272e+00, -7.86781e+00, -8.23192e+00],\n",
       "            [ 9.60838e-02, -2.20267e-01, -4.96593e-01,  ..., -5.02743e+00, -4.89260e+00, -6.91897e+00]],\n",
       " \n",
       "           [[-2.56901e-01, -2.44830e-01, -8.26820e-01,  ..., -8.11956e+00, -6.82304e+00, -7.41232e+00],\n",
       "            [-6.54104e-01,  1.20878e-01, -5.23452e-01,  ..., -1.19491e+01, -9.23878e+00, -8.54036e+00],\n",
       "            [-7.24554e-01, -5.02736e-02, -5.39922e-01,  ..., -1.17664e+01, -8.30877e+00, -8.05356e+00],\n",
       "            ...,\n",
       "            [ 2.71977e-01, -3.40064e-01, -1.86464e-01,  ..., -9.05405e+00, -6.25063e+00, -7.10211e+00],\n",
       "            [-1.40923e-01,  6.06483e-01, -7.46097e-01,  ..., -1.03899e+01, -8.29614e+00, -8.61360e+00],\n",
       "            [ 1.20650e-01, -2.80616e-01, -6.08180e-01,  ..., -6.25609e+00, -4.94420e+00, -7.46872e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.96421e-01, -6.21629e-02, -7.67842e-01,  ..., -6.08666e+00, -5.15821e+00, -7.09497e+00],\n",
       "            [ 8.21850e-01, -7.79213e-01, -6.31857e-01,  ..., -9.67955e+00, -7.87789e+00, -8.53824e+00],\n",
       "            [-1.27742e-01,  4.25492e-01, -4.10188e-01,  ..., -1.04194e+01, -7.91265e+00, -8.72065e+00],\n",
       "            ...,\n",
       "            [-3.94073e-01, -1.18016e-01, -6.86178e-01,  ..., -1.16111e+01, -8.96809e+00, -8.19850e+00],\n",
       "            [-6.66850e-01, -6.42969e-02, -3.47444e-01,  ..., -1.00462e+01, -7.89059e+00, -7.98691e+00],\n",
       "            [ 1.66660e-01, -5.78055e-01, -4.54455e-01,  ..., -6.88762e+00, -4.89413e+00, -7.24014e+00]],\n",
       " \n",
       "           [[-6.75319e-01,  8.13547e-02, -7.02954e-01,  ..., -5.89226e+00, -4.41401e+00, -6.85461e+00],\n",
       "            [-9.54965e-01,  1.05621e-01, -4.85575e-01,  ..., -9.91567e+00, -8.66195e+00, -8.40608e+00],\n",
       "            [ 4.76768e-02,  9.42831e-01,  4.33553e-02,  ..., -1.02490e+01, -8.32147e+00, -8.18778e+00],\n",
       "            ...,\n",
       "            [ 6.32594e-01,  1.15113e-01, -8.68662e-01,  ..., -1.15313e+01, -8.82642e+00, -8.85940e+00],\n",
       "            [ 8.08788e-01, -4.63330e-01, -4.33187e-01,  ..., -1.04264e+01, -7.82714e+00, -8.20156e+00],\n",
       "            [-8.88299e-02,  8.20136e-01, -2.76089e-01,  ..., -5.88826e+00, -4.87270e+00, -6.86476e+00]],\n",
       " \n",
       "           [[-2.15808e-01,  4.22767e-01, -1.74553e-01,  ..., -3.82294e+00, -3.94180e+00, -5.42703e+00],\n",
       "            [-1.21635e+00,  3.07135e-01, -4.83961e-03,  ..., -4.68567e+00, -4.86017e+00, -6.34192e+00],\n",
       "            [-6.71206e-02,  3.30766e-01,  2.07142e-01,  ..., -4.02695e+00, -4.71229e+00, -5.00468e+00],\n",
       "            ...,\n",
       "            [ 3.23773e-01,  7.75946e-01, -3.46639e-01,  ..., -5.93059e+00, -5.14683e+00, -6.83072e+00],\n",
       "            [ 2.20673e-01,  8.05406e-01, -1.34980e-01,  ..., -4.86114e+00, -5.43153e+00, -6.76441e+00],\n",
       "            [-1.11341e-01,  1.62591e-01, -5.27672e-02,  ..., -3.31917e+00, -4.06829e+00, -5.33628e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.56980e-01, -1.60930e-01, -6.06594e-01,  ..., -4.62826e+00, -1.85990e+00, -4.77936e+00],\n",
       "            [-3.88844e-01, -1.31099e-01, -3.36051e-01,  ..., -5.67309e+00, -3.01335e+00, -5.98691e+00],\n",
       "            [-5.03120e-01,  2.26705e-02, -5.40661e-01,  ..., -5.89764e+00, -4.22763e+00, -5.80133e+00],\n",
       "            ...,\n",
       "            [-3.97411e-01, -3.55976e-01, -1.74684e-01,  ..., -5.54950e+00, -4.03135e+00, -5.88270e+00],\n",
       "            [ 2.18803e-01, -6.44650e-01,  3.87803e-02,  ..., -5.53618e+00, -2.95747e+00, -5.68168e+00],\n",
       "            [ 1.96978e-02, -8.32281e-02, -3.97110e-01,  ..., -4.13821e+00, -2.37165e+00, -4.50097e+00]],\n",
       " \n",
       "           [[-4.69265e-01,  3.77661e-02, -9.33416e-01,  ..., -5.47055e+00, -4.28319e+00, -5.86433e+00],\n",
       "            [ 1.98315e-01, -7.50402e-01, -6.44417e-01,  ..., -7.10430e+00, -4.26897e+00, -7.06277e+00],\n",
       "            [ 4.69270e-01, -2.71538e-01, -5.60969e-01,  ..., -7.70977e+00, -4.45657e+00, -7.67100e+00],\n",
       "            ...,\n",
       "            [ 7.24788e-01,  9.13289e-02, -6.92392e-01,  ..., -6.89507e+00, -5.35624e+00, -6.85231e+00],\n",
       "            [ 6.79404e-01, -9.03883e-01, -6.00573e-01,  ..., -6.45955e+00, -4.58302e+00, -6.55330e+00],\n",
       "            [ 5.91324e-02, -3.07966e-01, -6.87622e-01,  ..., -4.53326e+00, -3.16620e+00, -4.74065e+00]],\n",
       " \n",
       "           [[-9.27686e-02, -4.35970e-01, -7.86723e-01,  ..., -6.22376e+00, -3.88801e+00, -5.70880e+00],\n",
       "            [-5.69414e-01,  9.59057e-03, -4.56796e-01,  ..., -8.54687e+00, -5.14144e+00, -7.53856e+00],\n",
       "            [-6.79464e-01, -1.97362e-01, -4.25779e-01,  ..., -8.54354e+00, -4.04835e+00, -7.42977e+00],\n",
       "            ...,\n",
       "            [ 2.78334e-01, -4.23337e-01, -2.65380e-01,  ..., -7.28902e+00, -2.45895e+00, -6.39896e+00],\n",
       "            [-1.30820e-01,  5.55337e-01, -6.92364e-01,  ..., -7.45119e+00, -4.78391e+00, -7.52191e+00],\n",
       "            [ 8.32029e-02, -3.84061e-01, -6.58098e-01,  ..., -5.33125e+00, -2.23498e+00, -5.65157e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.51886e-01, -1.89810e-01, -8.19604e-01,  ..., -5.33286e+00, -3.13860e+00, -4.78821e+00],\n",
       "            [ 8.64020e-01, -8.39708e-01, -6.43157e-01,  ..., -6.60470e+00, -4.37516e+00, -6.65722e+00],\n",
       "            [-6.25527e-02,  3.52172e-01, -3.68287e-01,  ..., -7.31488e+00, -3.72555e+00, -7.16098e+00],\n",
       "            ...,\n",
       "            [-3.60090e-01, -2.06571e-01, -5.65237e-01,  ..., -8.11784e+00, -4.98902e+00, -7.67828e+00],\n",
       "            [-6.49835e-01, -1.88209e-01, -3.63847e-01,  ..., -7.37570e+00, -3.69828e+00, -7.15522e+00],\n",
       "            [ 4.55839e-02, -6.17495e-01, -5.97206e-01,  ..., -5.36842e+00, -2.26768e+00, -5.41873e+00]],\n",
       " \n",
       "           [[-5.10753e-01, -5.09993e-02, -7.75230e-01,  ..., -5.09866e+00, -2.09601e+00, -4.85341e+00],\n",
       "            [-8.70055e-01,  1.49840e-02, -4.82842e-01,  ..., -6.36494e+00, -4.23290e+00, -6.38901e+00],\n",
       "            [ 1.04286e-01,  8.71486e-01, -7.35095e-02,  ..., -6.80268e+00, -3.66078e+00, -6.81703e+00],\n",
       "            ...,\n",
       "            [ 6.50943e-01,  1.03332e-01, -7.24566e-01,  ..., -7.24646e+00, -4.96183e+00, -7.40996e+00],\n",
       "            [ 8.00176e-01, -5.49073e-01, -4.15033e-01,  ..., -6.95256e+00, -4.18118e+00, -7.20831e+00],\n",
       "            [-2.15817e-01,  7.27997e-01, -4.35660e-01,  ..., -5.30751e+00, -2.75680e+00, -5.25540e+00]],\n",
       " \n",
       "           [[-1.30746e-01,  3.43234e-01, -5.73160e-01,  ..., -4.16249e+00, -2.88120e+00, -2.89621e+00],\n",
       "            [-1.17368e+00,  2.70090e-01, -3.14814e-01,  ..., -4.48188e+00, -3.49187e+00, -4.36516e+00],\n",
       "            [-7.12444e-02,  2.47548e-01, -1.06931e-01,  ..., -4.82979e+00, -3.50373e+00, -3.99628e+00],\n",
       "            ...,\n",
       "            [ 3.53521e-01,  5.98988e-01, -3.46746e-01,  ..., -5.00640e+00, -3.20755e+00, -5.11315e+00],\n",
       "            [ 2.11606e-01,  6.90252e-01, -3.06950e-01,  ..., -4.75140e+00, -3.93944e+00, -5.16873e+00],\n",
       "            [-2.08484e-01,  6.72733e-02, -4.39410e-01,  ..., -4.56471e+00, -3.45437e+00, -3.62267e+00]]]]], device='cuda:0'))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_model(torch.rand((1,3,384, 672)).cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f7913a66",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (model): Sequential(\n",
       "    (0): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (1): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (2): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (3): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (4): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (5): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (6): Concat()\n",
       "    (7): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (8): MP(\n",
       "      (m): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (9): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (10): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (11): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (12): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (13): Concat()\n",
       "    (14): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (15): MP(\n",
       "      (m): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (16): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (17): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (18): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (19): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (20): Concat()\n",
       "    (21): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (22): MP(\n",
       "      (m): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (23): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (24): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (25): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (26): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (27): Concat()\n",
       "    (28): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        1024, 512, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (29): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (30): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (31): SP(\n",
       "      (m): MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (32): SP(\n",
       "      (m): MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (33): SP(\n",
       "      (m): MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (34): Concat()\n",
       "    (35): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (36): Concat()\n",
       "    (37): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (38): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (39): Upsample(scale_factor=2.0, mode=nearest)\n",
       "    (40): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (41): Concat()\n",
       "    (42): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (43): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (44): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (45): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (46): Concat()\n",
       "    (47): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (48): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (49): Upsample(scale_factor=2.0, mode=nearest)\n",
       "    (50): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (51): Concat()\n",
       "    (52): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (53): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (54): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (55): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (56): Concat()\n",
       "    (57): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (58): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (59): Concat()\n",
       "    (60): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (61): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (62): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (63): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (64): Concat()\n",
       "    (65): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (66): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (67): Concat()\n",
       "    (68): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (69): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (70): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (71): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (72): Concat()\n",
       "    (73): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (74): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (75): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (76): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (77): IDetect(\n",
       "      (m): ModuleList(\n",
       "        (0): Conv2d(128, 45, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (1): Conv2d(256, 45, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (2): Conv2d(512, 45, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (ia): ModuleList(\n",
       "        (0): ImplicitA()\n",
       "        (1): ImplicitA()\n",
       "        (2): ImplicitA()\n",
       "      )\n",
       "      (im): ModuleList(\n",
       "        (0): ImplicitM()\n",
       "        (1): ImplicitM()\n",
       "        (2): ImplicitM()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6acc3670",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[[[[-4.13326e-01, -6.23371e-02,  5.75005e-01,  ..., -5.31502e+00, -5.16107e+00, -2.58336e+00],\n",
       "            [-1.75843e-01, -2.12433e-01,  8.34009e-01,  ..., -6.20939e+00, -5.19469e+00, -3.07918e+00],\n",
       "            [-4.22990e-01, -1.98272e-01,  8.13441e-01,  ..., -6.89601e+00, -5.63713e+00, -3.84209e+00],\n",
       "            ...,\n",
       "            [ 1.32598e-01,  5.98902e-02,  9.51860e-01,  ..., -7.06737e+00, -5.26251e+00, -4.33658e+00],\n",
       "            [-7.16299e-01, -4.58673e-02,  1.16375e+00,  ..., -7.49755e+00, -6.05372e+00, -5.10794e+00],\n",
       "            [ 3.38462e-02, -1.87414e-02,  8.11759e-01,  ..., -6.72266e+00, -6.11831e+00, -3.99034e+00]],\n",
       " \n",
       "           [[-2.16347e-01, -7.94555e-01,  5.66289e-01,  ..., -5.38766e+00, -5.03973e+00, -2.49367e+00],\n",
       "            [-3.46139e-01, -3.42304e-01,  9.74601e-01,  ..., -5.99159e+00, -4.85145e+00, -3.14609e+00],\n",
       "            [-7.45145e-01,  6.87071e-02,  9.42374e-01,  ..., -5.85254e+00, -4.92627e+00, -3.48064e+00],\n",
       "            ...,\n",
       "            [-6.90832e-01, -1.23752e-01,  1.04019e+00,  ..., -7.40004e+00, -5.59156e+00, -5.08472e+00],\n",
       "            [-3.98795e-01, -4.99748e-01,  8.81537e-01,  ..., -8.29504e+00, -6.56144e+00, -6.34111e+00],\n",
       "            [-1.76484e-01, -1.10627e+00,  7.15026e-01,  ..., -5.89737e+00, -5.29667e+00, -3.06933e+00]],\n",
       " \n",
       "           [[-1.69522e-01, -1.41861e-01,  4.89189e-01,  ..., -5.62718e+00, -4.67066e+00, -2.98102e+00],\n",
       "            [ 8.79851e-01,  1.85727e-01,  7.64724e-01,  ..., -6.99402e+00, -5.42635e+00, -4.25698e+00],\n",
       "            [-5.51406e-01,  9.52308e-02,  7.18871e-01,  ..., -7.60163e+00, -6.48608e+00, -6.26647e+00],\n",
       "            ...,\n",
       "            [-1.27266e+00, -6.66253e-01,  6.37504e-01,  ..., -7.26127e+00, -4.99074e+00, -5.58839e+00],\n",
       "            [-3.40734e-01, -9.85984e-01,  8.96658e-01,  ..., -7.08319e+00, -5.50628e+00, -4.75478e+00],\n",
       "            [ 1.36147e-01, -4.77219e-01,  5.73893e-01,  ..., -5.95032e+00, -5.31563e+00, -3.41680e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.75127e-01, -2.21594e-01,  5.32059e-01,  ..., -6.50069e+00, -5.46927e+00, -3.15628e+00],\n",
       "            [-6.18121e-01, -1.70262e-03,  6.67091e-01,  ..., -6.62617e+00, -5.04159e+00, -4.23259e+00],\n",
       "            [-7.47770e-01,  1.42828e-01,  9.03252e-01,  ..., -6.44738e+00, -4.15800e+00, -3.86356e+00],\n",
       "            ...,\n",
       "            [-1.25303e+00,  2.75804e-01,  7.60936e-01,  ..., -6.36797e+00, -4.26057e+00, -4.50545e+00],\n",
       "            [-5.19282e-02,  2.50867e-01,  5.50440e-01,  ..., -7.65614e+00, -5.54641e+00, -5.03222e+00],\n",
       "            [-1.02275e-01, -8.23122e-01,  5.31697e-01,  ..., -6.45691e+00, -5.78015e+00, -3.60641e+00]],\n",
       " \n",
       "           [[-6.19049e-01,  1.00231e-01,  4.16572e-01,  ..., -6.39538e+00, -5.36697e+00, -2.92152e+00],\n",
       "            [-5.89885e-01,  8.83334e-01,  6.36622e-01,  ..., -6.94260e+00, -5.49579e+00, -4.63012e+00],\n",
       "            [-8.78403e-01,  1.20517e+00,  8.78872e-01,  ..., -7.01354e+00, -5.46730e+00, -4.49490e+00],\n",
       "            ...,\n",
       "            [-1.34368e+00,  7.37194e-01,  4.95019e-01,  ..., -7.08247e+00, -5.38067e+00, -4.80377e+00],\n",
       "            [-6.54876e-01,  7.46498e-01,  5.43766e-01,  ..., -7.53969e+00, -5.77801e+00, -4.57719e+00],\n",
       "            [-6.15577e-01,  3.03322e-01,  6.55408e-01,  ..., -6.06812e+00, -5.27259e+00, -2.90809e+00]],\n",
       " \n",
       "           [[-6.64446e-01, -6.49181e-01,  5.97502e-01,  ..., -5.83543e+00, -5.26828e+00, -1.93361e+00],\n",
       "            [-5.89360e-01, -1.99928e-01,  7.85012e-01,  ..., -6.30711e+00, -5.70038e+00, -2.60049e+00],\n",
       "            [-7.16779e-01, -2.60577e-01,  8.77475e-01,  ..., -6.53440e+00, -5.68983e+00, -2.80189e+00],\n",
       "            ...,\n",
       "            [-9.68709e-01, -1.35809e-01,  6.57967e-01,  ..., -6.08348e+00, -5.42832e+00, -3.23978e+00],\n",
       "            [-6.59709e-01, -2.45439e-01,  6.54953e-01,  ..., -6.24352e+00, -5.73974e+00, -2.24688e+00],\n",
       "            [-6.50576e-01, -4.62507e-01,  5.50762e-01,  ..., -5.87213e+00, -5.17047e+00, -2.51659e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 3.15241e-02,  1.67111e-01,  5.70455e-01,  ..., -4.94009e+00, -5.01163e+00, -1.84756e+00],\n",
       "            [ 8.46715e-01, -1.23687e-01,  8.95811e-01,  ..., -6.22543e+00, -5.46682e+00, -1.97176e+00],\n",
       "            [ 3.33693e-01, -7.25198e-02,  6.81430e-01,  ..., -7.22741e+00, -5.90072e+00, -3.13352e+00],\n",
       "            ...,\n",
       "            [ 1.26796e+00,  3.40933e-01,  8.82812e-01,  ..., -7.73897e+00, -5.30650e+00, -3.01340e+00],\n",
       "            [ 5.31813e-01,  1.14184e-01,  1.05261e+00,  ..., -8.24829e+00, -6.44301e+00, -3.63851e+00],\n",
       "            [ 6.10001e-02,  1.72418e-01,  4.35136e-01,  ..., -6.97873e+00, -6.18018e+00, -2.90302e+00]],\n",
       " \n",
       "           [[ 1.12772e-02, -3.71459e-01,  6.10130e-01,  ..., -5.42351e+00, -4.92740e+00, -2.02399e+00],\n",
       "            [ 7.23210e-01, -1.55400e-01,  1.24525e+00,  ..., -6.26084e+00, -5.08654e+00, -2.89973e+00],\n",
       "            [-7.99489e-02,  6.85569e-01,  1.06438e+00,  ..., -6.14062e+00, -5.06400e+00, -3.25752e+00],\n",
       "            ...,\n",
       "            [ 6.95385e-01, -1.56455e-01,  1.06107e+00,  ..., -8.49025e+00, -5.81645e+00, -4.44979e+00],\n",
       "            [ 3.25241e-01, -3.85062e-01,  7.77698e-01,  ..., -9.81198e+00, -6.98146e+00, -5.74092e+00],\n",
       "            [ 3.00965e-01, -6.87397e-01,  5.08346e-01,  ..., -6.01987e+00, -5.28080e+00, -2.64167e+00]],\n",
       " \n",
       "           [[ 3.14099e-02,  1.89603e-01,  5.08301e-01,  ..., -5.52678e+00, -4.32298e+00, -2.80997e+00],\n",
       "            [ 1.83390e+00,  2.26865e-01,  7.90584e-01,  ..., -7.41339e+00, -5.60614e+00, -3.56491e+00],\n",
       "            [-5.18606e-01,  1.80995e-01,  5.29921e-01,  ..., -8.53779e+00, -6.82570e+00, -5.82524e+00],\n",
       "            ...,\n",
       "            [-5.85837e-01, -3.56575e-01,  6.53573e-01,  ..., -8.26155e+00, -5.17047e+00, -5.24155e+00],\n",
       "            [ 5.11615e-01, -6.41541e-01,  9.18563e-01,  ..., -8.00580e+00, -5.62340e+00, -4.18504e+00],\n",
       "            [ 2.57962e-01, -5.12527e-02,  3.54129e-01,  ..., -6.19470e+00, -5.08105e+00, -3.02417e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 1.07982e-01, -5.53778e-03,  5.64231e-01,  ..., -6.90756e+00, -5.53039e+00, -2.91890e+00],\n",
       "            [ 2.60366e-01,  1.66095e-01,  7.85530e-01,  ..., -7.19661e+00, -5.25316e+00, -4.29562e+00],\n",
       "            [ 9.78057e-01,  2.09775e-01,  1.14489e+00,  ..., -6.99659e+00, -4.51050e+00, -4.02218e+00],\n",
       "            ...,\n",
       "            [ 4.81205e-01,  2.35977e-01,  9.03587e-01,  ..., -7.07956e+00, -4.64272e+00, -4.25424e+00],\n",
       "            [ 1.16130e+00,  3.95851e-01,  4.14173e-01,  ..., -8.91073e+00, -5.76622e+00, -4.78688e+00],\n",
       "            [ 3.86165e-01, -3.07194e-01,  2.71970e-01,  ..., -7.17116e+00, -5.77716e+00, -3.34195e+00]],\n",
       " \n",
       "           [[ 2.18856e-01,  5.50570e-01,  5.27042e-01,  ..., -6.62805e+00, -5.14611e+00, -2.66948e+00],\n",
       "            [ 3.81243e-01,  1.17858e+00,  5.89297e-01,  ..., -7.86200e+00, -5.56699e+00, -4.26474e+00],\n",
       "            [ 4.40679e-01,  1.40129e+00,  8.30997e-01,  ..., -7.86126e+00, -5.46153e+00, -4.07662e+00],\n",
       "            ...,\n",
       "            [ 2.51116e-01,  8.27920e-01,  4.22686e-01,  ..., -8.25877e+00, -5.49661e+00, -4.58303e+00],\n",
       "            [ 6.01357e-01,  9.15833e-01,  4.80232e-01,  ..., -8.78634e+00, -5.89633e+00, -4.29897e+00],\n",
       "            [ 2.74621e-01,  6.65319e-01,  4.85792e-01,  ..., -6.44201e+00, -5.23809e+00, -2.79901e+00]],\n",
       " \n",
       "           [[ 3.76069e-01, -6.87711e-01,  6.55230e-01,  ..., -5.92175e+00, -5.06701e+00, -2.06170e+00],\n",
       "            [ 1.52795e-01, -2.34344e-01,  6.25847e-01,  ..., -6.95866e+00, -5.66122e+00, -2.48762e+00],\n",
       "            [ 4.43116e-01, -3.28291e-01,  6.93550e-01,  ..., -7.13827e+00, -5.63657e+00, -2.71700e+00],\n",
       "            ...,\n",
       "            [-2.62202e-01, -1.05257e-01,  5.25813e-01,  ..., -6.56651e+00, -5.31970e+00, -3.23756e+00],\n",
       "            [ 6.76957e-01, -2.10369e-01,  4.92186e-01,  ..., -7.01511e+00, -5.82348e+00, -2.25335e+00],\n",
       "            [ 1.86476e-01, -4.30283e-01,  3.21836e-01,  ..., -6.00715e+00, -5.17789e+00, -2.51570e+00]]],\n",
       " \n",
       " \n",
       "          [[[-1.59826e-02,  3.99077e-01,  2.07516e-02,  ..., -3.55377e+00, -4.87501e+00, -1.77531e+00],\n",
       "            [ 1.19874e-01, -1.98512e-02,  5.64315e-01,  ..., -4.75893e+00, -4.71322e+00, -1.94614e+00],\n",
       "            [-3.62436e-01, -1.31910e-02,  3.03685e-01,  ..., -5.73783e+00, -5.16576e+00, -3.12259e+00],\n",
       "            ...,\n",
       "            [ 6.79095e-01,  4.17015e-01,  4.35131e-01,  ..., -6.25494e+00, -4.37598e+00, -2.39291e+00],\n",
       "            [-6.50900e-02,  1.10051e-01,  4.78290e-01,  ..., -6.99741e+00, -5.42534e+00, -2.50072e+00],\n",
       "            [-2.25398e-01,  2.68510e-01,  8.50625e-02,  ..., -5.89355e+00, -5.80142e+00, -2.22436e+00]],\n",
       " \n",
       "           [[ 4.31748e-02,  8.91218e-02, -4.28416e-02,  ..., -4.25395e+00, -4.70145e+00, -2.34147e+00],\n",
       "            [-6.19578e-02,  1.62302e-01,  7.27785e-01,  ..., -5.22391e+00, -4.28238e+00, -3.82304e+00],\n",
       "            [-5.44352e-01,  1.32833e+00,  4.89184e-01,  ..., -5.31538e+00, -4.25520e+00, -3.96869e+00],\n",
       "            ...,\n",
       "            [-2.24084e-01,  1.30979e-02,  4.33939e-01,  ..., -7.83190e+00, -4.80156e+00, -4.27267e+00],\n",
       "            [-2.46105e-01, -2.54597e-01,  2.73480e-01,  ..., -9.13116e+00, -5.76216e+00, -4.70109e+00],\n",
       "            [-2.53768e-01, -2.76195e-01,  2.05780e-01,  ..., -5.16344e+00, -4.64446e+00, -1.87401e+00]],\n",
       " \n",
       "           [[ 1.93925e-02,  5.71791e-01, -1.05010e-01,  ..., -4.22785e+00, -4.36921e+00, -3.05285e+00],\n",
       "            [ 1.17526e+00,  3.42665e-01,  2.75091e-01,  ..., -6.40321e+00, -5.02255e+00, -3.38384e+00],\n",
       "            [-5.96325e-01,  2.48172e-01, -6.01764e-02,  ..., -7.58370e+00, -6.16070e+00, -5.29598e+00],\n",
       "            ...,\n",
       "            [-1.14165e+00,  2.81611e-02,  1.53633e-01,  ..., -7.95950e+00, -4.15345e+00, -5.11931e+00],\n",
       "            [-7.27785e-02, -2.37632e-01,  3.82246e-01,  ..., -7.57105e+00, -4.54904e+00, -4.23992e+00],\n",
       "            [-8.89839e-02,  3.61792e-01, -1.95352e-02,  ..., -5.42081e+00, -4.57197e+00, -2.49687e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.10565e-01,  2.66693e-01, -4.42389e-02,  ..., -5.94173e+00, -4.96589e+00, -2.74422e+00],\n",
       "            [-1.84236e-01,  3.86925e-01,  3.95337e-01,  ..., -6.37514e+00, -4.29142e+00, -4.36154e+00],\n",
       "            [ 5.60648e-02,  4.28676e-01,  6.06334e-01,  ..., -6.19639e+00, -3.52790e+00, -4.05501e+00],\n",
       "            ...,\n",
       "            [-3.56128e-01,  4.55358e-01,  4.51467e-01,  ..., -6.34572e+00, -3.52759e+00, -4.60054e+00],\n",
       "            [ 3.52770e-01,  7.00180e-01, -2.59291e-02,  ..., -8.31957e+00, -4.73096e+00, -5.37422e+00],\n",
       "            [ 3.90062e-02,  1.88525e-01, -1.50090e-01,  ..., -6.61097e+00, -5.10333e+00, -3.01932e+00]],\n",
       " \n",
       "           [[-6.24542e-02,  8.79207e-01, -4.95053e-02,  ..., -5.45732e+00, -4.66802e+00, -2.91482e+00],\n",
       "            [-2.68343e-01,  1.37272e+00,  1.28359e-01,  ..., -6.77138e+00, -4.76336e+00, -4.53291e+00],\n",
       "            [-2.18886e-01,  1.51231e+00,  2.53387e-01,  ..., -6.81574e+00, -4.68664e+00, -4.68831e+00],\n",
       "            ...,\n",
       "            [-5.84782e-01,  1.05083e+00, -1.06287e-01,  ..., -7.41213e+00, -4.64073e+00, -5.47576e+00],\n",
       "            [-8.82299e-02,  1.01793e+00,  1.81268e-03,  ..., -7.97613e+00, -5.03264e+00, -5.21954e+00],\n",
       "            [-3.70566e-01,  9.55625e-01,  7.09876e-02,  ..., -5.72601e+00, -4.63692e+00, -2.94823e+00]],\n",
       " \n",
       "           [[ 9.19310e-02, -4.79952e-01,  1.10158e-01,  ..., -4.95794e+00, -4.81186e+00, -1.58768e+00],\n",
       "            [-3.99809e-01, -1.52557e-01,  1.53243e-01,  ..., -6.24999e+00, -5.13912e+00, -2.47357e+00],\n",
       "            [-1.61226e-01, -2.14606e-01,  1.44195e-01,  ..., -6.44162e+00, -5.16575e+00, -2.93264e+00],\n",
       "            ...,\n",
       "            [-6.58835e-01,  6.53667e-02, -9.96307e-03,  ..., -5.91151e+00, -5.00694e+00, -3.21644e+00],\n",
       "            [-7.34063e-02, -3.13016e-02,  1.48874e-02,  ..., -6.46881e+00, -5.25530e+00, -2.65170e+00],\n",
       "            [-3.49927e-01, -2.67646e-01, -7.64554e-02,  ..., -5.16526e+00, -4.96780e+00, -2.21488e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-4.27043e-01, -2.28936e-01,  1.27014e-01,  ..., -6.62655e+00, -2.47402e+00, -5.66630e+00],\n",
       "            [ 2.78863e-02, -6.08720e-01,  7.81907e-01,  ..., -7.23107e+00, -2.97633e+00, -5.48184e+00],\n",
       "            [-4.89116e-01, -5.19087e-01,  3.25076e-01,  ..., -8.34523e+00, -4.12188e+00, -5.90691e+00],\n",
       "            ...,\n",
       "            [-4.05199e-01, -2.47564e-01,  6.75558e-01,  ..., -7.53310e+00, -4.04751e+00, -5.49954e+00],\n",
       "            [ 2.86968e-01, -1.88718e-01,  8.68360e-01,  ..., -7.89050e+00, -3.51388e+00, -5.64862e+00],\n",
       "            [-1.92176e-01, -2.28015e-01,  7.51170e-01,  ..., -7.13643e+00, -3.07206e+00, -6.08342e+00]],\n",
       " \n",
       "           [[-2.08572e-01, -4.00339e-01,  1.97742e-01,  ..., -7.31968e+00, -3.75648e+00, -5.58432e+00],\n",
       "            [-6.66483e-01, -2.70773e-01,  5.82036e-01,  ..., -8.91705e+00, -3.91366e+00, -5.94962e+00],\n",
       "            [-7.16378e-02,  9.29686e-01,  6.62453e-01,  ..., -9.28227e+00, -4.54465e+00, -5.79770e+00],\n",
       "            ...,\n",
       "            [ 9.01068e-03,  7.60649e-01,  5.63273e-01,  ..., -9.39951e+00, -4.35840e+00, -5.58487e+00],\n",
       "            [-3.63491e-02,  7.25811e-01,  7.03924e-01,  ..., -9.72475e+00, -4.12710e+00, -5.64075e+00],\n",
       "            [ 3.20985e-01, -4.33883e-01,  2.90493e-01,  ..., -7.53156e+00, -3.07337e+00, -5.69699e+00]],\n",
       " \n",
       "           [[-5.27646e-01,  3.04164e-02, -3.99002e-02,  ..., -8.11126e+00, -4.31070e+00, -6.03176e+00],\n",
       "            [ 2.54693e-01,  7.08190e-01,  5.63379e-01,  ..., -9.42173e+00, -4.42820e+00, -5.74583e+00],\n",
       "            [ 2.30708e-01,  5.19283e-01,  4.19696e-01,  ..., -1.09680e+01, -5.69371e+00, -6.40645e+00],\n",
       "            ...,\n",
       "            [-7.63886e-02,  4.26026e-01,  4.96687e-01,  ..., -1.03172e+01, -4.61996e+00, -6.13749e+00],\n",
       "            [-3.32792e-01, -2.13901e-01,  5.40802e-01,  ..., -9.96149e+00, -4.43899e+00, -6.06601e+00],\n",
       "            [ 4.54634e-01,  7.30961e-02,  3.10645e-02,  ..., -9.52047e+00, -3.92313e+00, -6.79200e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.53254e-01,  5.97373e-01,  1.08003e-01,  ..., -8.58026e+00, -4.17475e+00, -5.94147e+00],\n",
       "            [ 3.49065e-01,  7.04439e-01,  6.61204e-01,  ..., -9.74512e+00, -4.81629e+00, -5.95154e+00],\n",
       "            [ 4.06500e-01,  8.79422e-02,  4.22821e-01,  ..., -9.92855e+00, -4.79600e+00, -5.72047e+00],\n",
       "            ...,\n",
       "            [ 4.66373e-01,  4.94262e-02,  4.46154e-01,  ..., -9.05642e+00, -4.39344e+00, -5.21804e+00],\n",
       "            [ 2.49975e-02, -9.65264e-01,  3.84365e-01,  ..., -9.53293e+00, -4.76976e+00, -5.75310e+00],\n",
       "            [ 4.41018e-01, -4.69861e-01,  6.98987e-02,  ..., -8.55857e+00, -3.43010e+00, -6.13263e+00]],\n",
       " \n",
       "           [[-4.27624e-01,  7.47981e-01,  1.81499e-01,  ..., -8.14956e+00, -4.58857e+00, -6.33422e+00],\n",
       "            [ 1.12416e-01,  1.28082e+00,  6.48242e-01,  ..., -8.96917e+00, -4.38030e+00, -5.91767e+00],\n",
       "            [ 3.44287e-01,  3.93331e-01,  3.94820e-01,  ..., -9.68147e+00, -4.31018e+00, -5.90271e+00],\n",
       "            ...,\n",
       "            [ 3.91448e-01,  6.62130e-01,  6.18897e-01,  ..., -9.09248e+00, -3.73604e+00, -5.75219e+00],\n",
       "            [-2.66316e-01,  1.52043e+00,  6.43627e-01,  ..., -9.05133e+00, -4.84906e+00, -6.03669e+00],\n",
       "            [ 3.01853e-01,  3.07736e-01,  1.93088e-01,  ..., -7.93199e+00, -3.07443e+00, -5.97177e+00]],\n",
       " \n",
       "           [[-3.11391e-01, -2.10427e-01,  2.69750e-01,  ..., -7.21454e+00, -3.72725e+00, -6.55247e+00],\n",
       "            [-3.63095e-01,  2.17890e-01,  6.00753e-01,  ..., -7.98768e+00, -4.23417e+00, -6.42905e+00],\n",
       "            [-2.98589e-01,  3.69496e-01,  4.03433e-01,  ..., -8.62938e+00, -4.21147e+00, -6.55268e+00],\n",
       "            ...,\n",
       "            [ 5.98737e-01,  7.15189e-02,  4.28356e-01,  ..., -8.50481e+00, -4.17908e+00, -6.66284e+00],\n",
       "            [-3.35211e-01,  1.73768e-01,  3.19881e-01,  ..., -8.06072e+00, -4.97071e+00, -6.50183e+00],\n",
       "            [ 2.51596e-01, -1.51395e-01,  2.62071e-01,  ..., -6.50311e+00, -3.59668e+00, -6.32020e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.66149e-01, -3.55863e-01, -2.53377e-01,  ..., -5.31219e+00, -2.28949e+00, -4.31402e+00],\n",
       "            [ 8.49542e-02, -7.22691e-01,  2.63838e-01,  ..., -5.19918e+00, -2.54136e+00, -4.81432e+00],\n",
       "            [-4.78971e-01, -6.06324e-01, -8.21602e-03,  ..., -5.48416e+00, -2.79422e+00, -4.38353e+00],\n",
       "            ...,\n",
       "            [-4.18691e-01, -3.51567e-01,  3.00135e-01,  ..., -5.04182e+00, -2.96468e+00, -4.41903e+00],\n",
       "            [ 3.92157e-01, -3.14557e-01,  4.46631e-01,  ..., -4.99819e+00, -2.81963e+00, -4.91282e+00],\n",
       "            [-3.14018e-01, -3.20499e-01,  1.38596e-01,  ..., -5.37607e+00, -2.58969e+00, -4.17664e+00]],\n",
       " \n",
       "           [[-9.76925e-02, -5.04887e-01, -2.38758e-01,  ..., -5.10996e+00, -3.38919e+00, -4.89736e+00],\n",
       "            [-7.48075e-01, -3.66053e-01,  4.62180e-02,  ..., -5.03499e+00, -3.21101e+00, -6.56734e+00],\n",
       "            [-1.08449e-01,  7.59888e-01,  1.82728e-01,  ..., -4.74289e+00, -3.23599e+00, -7.36696e+00],\n",
       "            ...,\n",
       "            [-6.08094e-03,  4.46010e-01,  4.62428e-02,  ..., -4.87335e+00, -2.90537e+00, -6.60880e+00],\n",
       "            [-3.18424e-02,  3.75379e-01,  1.34648e-01,  ..., -4.72520e+00, -2.49895e+00, -7.07473e+00],\n",
       "            [ 2.72556e-01, -6.55736e-01, -2.49891e-01,  ..., -5.36613e+00, -2.77365e+00, -5.21884e+00]],\n",
       " \n",
       "           [[-4.97228e-01, -7.38967e-02, -4.84245e-01,  ..., -5.50693e+00, -3.48428e+00, -4.92068e+00],\n",
       "            [ 1.78630e-01,  4.37279e-01,  3.46278e-02,  ..., -4.92661e+00, -3.30978e+00, -7.60202e+00],\n",
       "            [ 2.44477e-01,  3.78887e-01, -1.34439e-01,  ..., -5.23149e+00, -3.47120e+00, -8.53596e+00],\n",
       "            ...,\n",
       "            [-1.17282e-01,  2.21239e-01,  8.48707e-03,  ..., -5.11251e+00, -2.67544e+00, -6.78402e+00],\n",
       "            [-3.15157e-01, -2.91053e-01,  4.57756e-02,  ..., -5.10145e+00, -2.39347e+00, -6.04708e+00],\n",
       "            [ 4.50065e-01, -1.80091e-02, -4.87821e-01,  ..., -5.87374e+00, -2.19778e+00, -4.98028e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.35399e-01,  3.97269e-01, -3.84592e-01,  ..., -5.16295e+00, -3.29046e+00, -5.96878e+00],\n",
       "            [ 2.71769e-01,  6.05041e-01,  1.22100e-01,  ..., -4.92213e+00, -3.37502e+00, -8.03218e+00],\n",
       "            [ 3.24009e-01, -1.30230e-03, -2.42602e-02,  ..., -4.83585e+00, -2.85593e+00, -6.45356e+00],\n",
       "            ...,\n",
       "            [ 4.36083e-01, -1.14287e-01,  2.17725e-02,  ..., -4.54564e+00, -2.62989e+00, -7.02950e+00],\n",
       "            [ 7.83474e-02, -9.44465e-01, -7.96190e-03,  ..., -4.87005e+00, -2.87535e+00, -6.40160e+00],\n",
       "            [ 4.11584e-01, -5.26995e-01, -4.47351e-01,  ..., -5.33673e+00, -2.50511e+00, -5.99063e+00]],\n",
       " \n",
       "           [[-3.21408e-01,  7.19655e-01, -2.87782e-01,  ..., -5.27208e+00, -3.88372e+00, -5.01728e+00],\n",
       "            [ 1.90975e-02,  1.14754e+00,  1.32871e-01,  ..., -4.66875e+00, -3.46234e+00, -6.42581e+00],\n",
       "            [ 2.93530e-01,  2.83724e-01, -5.21029e-02,  ..., -4.70770e+00, -2.84567e+00, -6.38405e+00],\n",
       "            ...,\n",
       "            [ 4.10298e-01,  5.20079e-01,  9.63566e-02,  ..., -4.68001e+00, -3.00085e+00, -6.95945e+00],\n",
       "            [-2.52275e-01,  1.44008e+00,  1.98741e-01,  ..., -4.83724e+00, -3.77584e+00, -6.17653e+00],\n",
       "            [ 2.32832e-01,  2.11930e-01, -2.53558e-01,  ..., -5.12264e+00, -3.05741e+00, -6.33520e+00]],\n",
       " \n",
       "           [[-6.52034e-02, -1.66615e-01, -7.40748e-02,  ..., -5.49836e+00, -3.66220e+00, -3.50108e+00],\n",
       "            [-5.66149e-01,  1.37025e-01,  2.24628e-01,  ..., -5.34865e+00, -3.37382e+00, -4.23508e+00],\n",
       "            [-2.84079e-01,  3.07764e-01,  2.69451e-02,  ..., -5.51947e+00, -2.92877e+00, -4.36461e+00],\n",
       "            ...,\n",
       "            [ 7.00118e-01,  1.13322e-02,  2.49749e-02,  ..., -5.47352e+00, -3.47267e+00, -4.29762e+00],\n",
       "            [-2.34549e-01,  1.21807e-01,  1.13289e-01,  ..., -5.41964e+00, -4.08786e+00, -3.96526e+00],\n",
       "            [ 1.19993e-01, -1.96033e-01, -1.39221e-01,  ..., -5.54539e+00, -3.38369e+00, -3.68083e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.56872e-01, -1.84134e-01, -7.07281e-02,  ..., -5.39871e+00, -1.57926e+00, -4.05615e+00],\n",
       "            [ 8.38366e-02, -5.69163e-01,  5.87685e-01,  ..., -5.78874e+00, -2.01577e+00, -4.55516e+00],\n",
       "            [-4.56680e-01, -4.93161e-01,  2.69589e-01,  ..., -6.46818e+00, -2.93132e+00, -4.65104e+00],\n",
       "            ...,\n",
       "            [-3.99066e-01, -1.90444e-01,  6.61394e-01,  ..., -6.05003e+00, -3.05334e+00, -5.06528e+00],\n",
       "            [ 4.01199e-01, -1.41988e-01,  8.17730e-01,  ..., -6.19070e+00, -2.65790e+00, -5.07317e+00],\n",
       "            [-3.08035e-01, -1.82732e-01,  3.86149e-01,  ..., -5.66050e+00, -2.36756e+00, -3.80004e+00]],\n",
       " \n",
       "           [[-8.20363e-02, -4.21720e-01, -5.91869e-02,  ..., -5.70241e+00, -2.94912e+00, -4.20227e+00],\n",
       "            [-7.18715e-01, -3.32214e-01,  3.08806e-01,  ..., -7.10107e+00, -2.90653e+00, -6.77615e+00],\n",
       "            [-8.35778e-02,  8.39382e-01,  5.07340e-01,  ..., -7.30297e+00, -3.28614e+00, -7.77045e+00],\n",
       "            ...,\n",
       "            [ 3.18488e-02,  4.76121e-01,  3.61117e-01,  ..., -7.49254e+00, -2.73434e+00, -7.66200e+00],\n",
       "            [-5.77847e-03,  4.24098e-01,  4.61822e-01,  ..., -8.11854e+00, -2.65118e+00, -7.12609e+00],\n",
       "            [ 2.71308e-01, -5.57754e-01,  3.32192e-04,  ..., -6.15201e+00, -1.87505e+00, -5.20145e+00]],\n",
       " \n",
       "           [[-4.78144e-01,  8.52830e-03, -3.50140e-01,  ..., -6.16777e+00, -3.20274e+00, -4.05842e+00],\n",
       "            [ 2.07452e-01,  5.08260e-01,  2.86360e-01,  ..., -7.52394e+00, -3.09190e+00, -7.31601e+00],\n",
       "            [ 2.70671e-01,  4.60754e-01,  5.55834e-02,  ..., -8.66453e+00, -3.95098e+00, -8.90764e+00],\n",
       "            ...,\n",
       "            [-7.06554e-02,  2.88400e-01,  2.38410e-01,  ..., -8.56512e+00, -3.01093e+00, -7.80508e+00],\n",
       "            [-2.77887e-01, -2.49045e-01,  3.10279e-01,  ..., -8.59830e+00, -3.02401e+00, -8.00744e+00],\n",
       "            [ 4.52244e-01,  7.19065e-02, -2.94486e-01,  ..., -8.00632e+00, -2.51855e+00, -6.39020e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.13202e-01,  5.00746e-01, -2.54691e-01,  ..., -6.53701e+00, -3.03265e+00, -5.63204e+00],\n",
       "            [ 2.97613e-01,  6.66998e-01,  3.81439e-01,  ..., -8.04412e+00, -3.46514e+00, -7.74724e+00],\n",
       "            [ 3.45335e-01,  7.76238e-03,  2.77093e-01,  ..., -8.06992e+00, -3.12659e+00, -8.62766e+00],\n",
       "            ...,\n",
       "            [ 4.60836e-01, -9.11833e-02,  3.13988e-01,  ..., -7.62775e+00, -3.05795e+00, -8.58338e+00],\n",
       "            [ 1.06847e-01, -9.33932e-01,  2.57589e-01,  ..., -7.98656e+00, -3.76450e+00, -7.28438e+00],\n",
       "            [ 4.09218e-01, -4.96128e-01, -2.35114e-01,  ..., -7.17828e+00, -2.38698e+00, -6.29811e+00]],\n",
       " \n",
       "           [[-3.03986e-01,  8.09566e-01, -1.51413e-01,  ..., -6.16980e+00, -3.72078e+00, -4.59075e+00],\n",
       "            [ 2.82883e-02,  1.25756e+00,  3.93017e-01,  ..., -7.20325e+00, -3.47232e+00, -6.48701e+00],\n",
       "            [ 3.13729e-01,  3.25713e-01,  2.01141e-01,  ..., -7.96193e+00, -3.14965e+00, -7.26748e+00],\n",
       "            ...,\n",
       "            [ 4.18658e-01,  6.15943e-01,  3.54678e-01,  ..., -7.59254e+00, -2.90961e+00, -7.31787e+00],\n",
       "            [-2.39998e-01,  1.57241e+00,  4.92977e-01,  ..., -7.08510e+00, -3.90598e+00, -6.89802e+00],\n",
       "            [ 2.29018e-01,  3.55281e-01, -4.62747e-02,  ..., -6.82579e+00, -2.57963e+00, -5.17162e+00]],\n",
       " \n",
       "           [[-6.08647e-02, -1.77726e-01,  9.18569e-02,  ..., -5.60430e+00, -3.06597e+00, -2.67894e+00],\n",
       "            [-5.61206e-01,  1.13175e-01,  5.64568e-01,  ..., -6.12920e+00, -3.17475e+00, -3.45929e+00],\n",
       "            [-2.62180e-01,  2.61088e-01,  3.13622e-01,  ..., -6.83623e+00, -2.97462e+00, -4.04833e+00],\n",
       "            ...,\n",
       "            [ 6.93735e-01,  3.17566e-02,  2.64114e-01,  ..., -6.37598e+00, -3.31613e+00, -3.99989e+00],\n",
       "            [-2.27328e-01,  1.40650e-01,  4.53759e-01,  ..., -5.91029e+00, -3.97102e+00, -3.58434e+00],\n",
       "            [ 1.15020e-01, -1.94890e-01,  8.52809e-02,  ..., -5.35755e+00, -3.07325e+00, -2.45637e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-2.69908e-01,  2.25082e-01, -4.36437e-01,  ..., -5.07561e+00, -2.79527e+00, -3.32329e+00],\n",
       "            [-1.01081e+00, -4.84078e-01, -3.25506e-01,  ..., -5.53013e+00, -4.94407e+00, -2.94887e+00],\n",
       "            [ 1.55792e-01, -8.37627e-02, -4.84538e-01,  ..., -5.78167e+00, -5.73637e+00, -2.38510e+00],\n",
       "            ...,\n",
       "            [-9.61858e-02, -2.03941e-01, -3.00146e-01,  ..., -5.72660e+00, -5.63150e+00, -2.52486e+00],\n",
       "            [-4.09207e-02, -3.77281e-01, -4.07168e-01,  ..., -5.54893e+00, -5.24134e+00, -2.43417e+00],\n",
       "            [-7.19697e-02, -3.92714e-01, -1.32234e-01,  ..., -4.85605e+00, -2.98782e+00, -3.20555e+00]],\n",
       " \n",
       "           [[-7.48618e-01,  1.83701e-01, -8.58696e-01,  ..., -5.63096e+00, -4.13363e+00, -2.93145e+00],\n",
       "            [-1.98816e-01,  3.90269e-01, -8.49058e-01,  ..., -6.20742e+00, -6.39050e+00, -3.32109e+00],\n",
       "            [-7.07388e-02, -1.26051e-01, -6.56484e-01,  ..., -6.33459e+00, -7.53417e+00, -3.50814e+00],\n",
       "            ...,\n",
       "            [-7.60557e-01,  2.90967e-02, -8.14514e-01,  ..., -6.86741e+00, -8.50526e+00, -3.80182e+00],\n",
       "            [-1.01901e-01, -7.76213e-01, -5.24693e-01,  ..., -6.08733e+00, -7.17042e+00, -4.19816e+00],\n",
       "            [ 5.59415e-01,  2.33588e-01, -7.42672e-01,  ..., -6.09614e+00, -4.66380e+00, -3.43726e+00]],\n",
       " \n",
       "           [[-7.52911e-01,  1.77516e-01, -8.83402e-01,  ..., -5.64080e+00, -4.65287e+00, -2.35689e+00],\n",
       "            [-4.01591e-01,  1.93816e-01, -6.55977e-01,  ..., -6.46182e+00, -7.37380e+00, -3.65443e+00],\n",
       "            [ 4.61897e-01, -9.26231e-03, -4.59768e-01,  ..., -6.56443e+00, -8.48446e+00, -3.87300e+00],\n",
       "            ...,\n",
       "            [ 4.23227e-01, -9.04017e-02, -3.25530e-01,  ..., -6.05072e+00, -7.32842e+00, -4.16411e+00],\n",
       "            [-1.08190e+00,  2.70784e-01, -4.90661e-02,  ..., -5.85786e+00, -7.28135e+00, -4.78578e+00],\n",
       "            [-4.07758e-01, -1.30359e-01, -3.94476e-01,  ..., -5.13960e+00, -4.75030e+00, -3.17788e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-8.30927e-01,  2.34828e-01, -6.82914e-01,  ..., -5.39463e+00, -3.04103e+00, -2.91664e+00],\n",
       "            [ 4.22100e-01,  2.59706e-01, -4.91967e-01,  ..., -6.05952e+00, -6.07994e+00, -3.55859e+00],\n",
       "            [ 6.35116e-01, -8.48418e-01, -2.17575e-01,  ..., -6.31480e+00, -7.82932e+00, -3.65326e+00],\n",
       "            ...,\n",
       "            [-4.18881e-02,  7.31478e-02, -6.59285e-01,  ..., -6.38986e+00, -6.18098e+00, -3.44553e+00],\n",
       "            [-6.03201e-01, -2.27702e-03, -5.10331e-01,  ..., -5.54348e+00, -5.49899e+00, -4.13689e+00],\n",
       "            [ 2.86992e-01,  7.21157e-01, -5.59591e-01,  ..., -5.39104e+00, -3.42258e+00, -3.84392e+00]],\n",
       " \n",
       "           [[-6.36008e-01,  2.09460e-01, -7.14214e-01,  ..., -5.60297e+00, -3.99867e+00, -2.71116e+00],\n",
       "            [ 5.18625e-01, -4.94798e-01, -3.22848e-01,  ..., -5.96412e+00, -6.19714e+00, -3.82780e+00],\n",
       "            [ 8.67014e-01,  6.46808e-01, -2.83372e-01,  ..., -6.45473e+00, -7.40285e+00, -3.49671e+00],\n",
       "            ...,\n",
       "            [ 5.84057e-01, -6.58916e-01, -3.58345e-01,  ..., -5.54584e+00, -4.84633e+00, -4.12377e+00],\n",
       "            [-8.85281e-01,  2.34269e-02, -3.12962e-01,  ..., -5.43973e+00, -6.43099e+00, -4.64338e+00],\n",
       "            [ 6.20489e-01, -6.46603e-01, -6.87335e-01,  ..., -5.47645e+00, -2.89264e+00, -3.81506e+00]],\n",
       " \n",
       "           [[-1.09634e-01,  2.58209e-01, -2.08308e-01,  ..., -4.94642e+00, -4.20564e+00, -2.34099e+00],\n",
       "            [-4.93674e-01,  4.98183e-01, -8.90302e-02,  ..., -4.82046e+00, -4.82705e+00, -2.51709e+00],\n",
       "            [-4.12998e-01,  9.71582e-01,  1.29912e-01,  ..., -4.82937e+00, -5.34807e+00, -2.43066e+00],\n",
       "            ...,\n",
       "            [-2.49007e-01,  1.40935e-01, -1.63150e-01,  ..., -5.00556e+00, -4.76950e+00, -2.49446e+00],\n",
       "            [ 8.35314e-01,  4.05906e-01,  6.25549e-02,  ..., -4.82711e+00, -4.05742e+00, -2.99971e+00],\n",
       "            [-1.20250e-01,  2.45645e-01, -7.77278e-02,  ..., -4.77998e+00, -3.51328e+00, -3.02946e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.75325e-01,  2.33083e-01, -3.53039e-01,  ..., -4.75014e+00, -4.15715e+00, -6.36333e+00],\n",
       "            [-1.00311e+00, -4.44016e-01, -3.17348e-01,  ..., -7.20536e+00, -6.06815e+00, -7.77317e+00],\n",
       "            [ 1.58524e-01, -4.85094e-02, -5.28077e-01,  ..., -7.52101e+00, -7.12138e+00, -7.75193e+00],\n",
       "            ...,\n",
       "            [-8.46773e-02, -1.61781e-01, -3.19679e-01,  ..., -7.12107e+00, -7.16640e+00, -7.73881e+00],\n",
       "            [-3.48808e-02, -3.53080e-01, -4.11854e-01,  ..., -6.84751e+00, -6.73437e+00, -7.50593e+00],\n",
       "            [-6.03372e-02, -3.87564e-01, -4.50865e-02,  ..., -3.30697e+00, -4.00124e+00, -6.16820e+00]],\n",
       " \n",
       "           [[-7.53137e-01,  1.82594e-01, -8.15844e-01,  ..., -6.19678e+00, -5.03592e+00, -7.19631e+00],\n",
       "            [-2.01907e-01,  3.99348e-01, -8.84051e-01,  ..., -1.02563e+01, -8.07117e+00, -8.51642e+00],\n",
       "            [-6.63034e-02, -9.54183e-02, -7.21474e-01,  ..., -1.03900e+01, -8.30382e+00, -8.10066e+00],\n",
       "            ...,\n",
       "            [-7.61623e-01,  3.91802e-02, -8.10971e-01,  ..., -1.22347e+01, -9.43403e+00, -9.09008e+00],\n",
       "            [-8.93855e-02, -7.54257e-01, -5.27228e-01,  ..., -9.00068e+00, -7.42014e+00, -8.23446e+00],\n",
       "            [ 5.73007e-01,  2.59415e-01, -7.03650e-01,  ..., -7.19263e+00, -5.86184e+00, -8.16364e+00]],\n",
       " \n",
       "           [[-7.62086e-01,  1.96483e-01, -8.62482e-01,  ..., -5.67080e+00, -5.39669e+00, -6.97443e+00],\n",
       "            [-4.07175e-01,  2.24230e-01, -6.90435e-01,  ..., -1.00217e+01, -8.57190e+00, -8.25103e+00],\n",
       "            [ 4.68108e-01,  2.73532e-02, -5.16728e-01,  ..., -1.14676e+01, -9.17939e+00, -8.33405e+00],\n",
       "            ...,\n",
       "            [ 4.19662e-01, -5.88815e-02, -3.50271e-01,  ..., -1.20050e+01, -8.83029e+00, -8.36283e+00],\n",
       "            [-1.07189e+00,  3.19418e-01,  5.24117e-02,  ..., -1.16412e+01, -8.33558e+00, -8.51247e+00],\n",
       "            [-3.95082e-01, -1.00751e-01, -3.73488e-01,  ..., -7.16712e+00, -6.04441e+00, -7.19840e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-8.39715e-01,  2.54699e-01, -6.72550e-01,  ..., -5.68373e+00, -4.87007e+00, -6.77587e+00],\n",
       "            [ 4.30944e-01,  2.98732e-01, -5.30955e-01,  ..., -1.03215e+01, -8.34842e+00, -8.15411e+00],\n",
       "            [ 6.43867e-01, -8.14644e-01, -2.79034e-01,  ..., -1.12087e+01, -9.46514e+00, -8.47123e+00],\n",
       "            ...,\n",
       "            [-3.93242e-02,  1.00220e-01, -7.31252e-01,  ..., -1.19467e+01, -8.30040e+00, -8.24647e+00],\n",
       "            [-5.91361e-01,  3.95540e-02, -5.77408e-01,  ..., -9.78092e+00, -6.93335e+00, -8.01912e+00],\n",
       "            [ 2.97492e-01,  7.65984e-01, -5.42973e-01,  ..., -7.14495e+00, -5.07616e+00, -7.79304e+00]],\n",
       " \n",
       "           [[-6.45194e-01,  2.23212e-01, -6.88823e-01,  ..., -6.28409e+00, -5.54723e+00, -7.03631e+00],\n",
       "            [ 5.18861e-01, -4.62516e-01, -2.89195e-01,  ..., -9.26187e+00, -8.09364e+00, -7.82651e+00],\n",
       "            [ 8.78749e-01,  6.95139e-01, -2.92990e-01,  ..., -1.05942e+01, -8.55832e+00, -8.22095e+00],\n",
       "            ...,\n",
       "            [ 5.89746e-01, -6.31788e-01, -4.00456e-01,  ..., -9.90013e+00, -7.36606e+00, -7.82909e+00],\n",
       "            [-8.71753e-01,  3.77831e-02, -2.92938e-01,  ..., -1.03046e+01, -8.19001e+00, -8.15463e+00],\n",
       "            [ 6.38880e-01, -6.63496e-01, -6.48966e-01,  ..., -6.09863e+00, -4.32731e+00, -7.39224e+00]],\n",
       " \n",
       "           [[-1.16607e-01,  2.46199e-01, -1.31162e-01,  ..., -4.29277e+00, -5.12158e+00, -5.42872e+00],\n",
       "            [-4.87220e-01,  4.96359e-01, -2.76936e-02,  ..., -5.01224e+00, -5.77913e+00, -6.23203e+00],\n",
       "            [-4.01414e-01,  9.71506e-01,  1.98478e-01,  ..., -5.15265e+00, -5.38268e+00, -6.55797e+00],\n",
       "            ...,\n",
       "            [-2.39561e-01,  1.33340e-01, -1.82304e-01,  ..., -7.12026e+00, -6.56846e+00, -6.71360e+00],\n",
       "            [ 8.46440e-01,  3.67794e-01,  1.22073e-01,  ..., -5.30289e+00, -5.23495e+00, -6.22452e+00],\n",
       "            [-1.06038e-01,  2.11365e-01,  4.10145e-02,  ..., -3.45334e+00, -3.85393e+00, -5.49510e+00]]],\n",
       " \n",
       " \n",
       "          [[[-1.49750e-01,  1.53865e-01, -5.73252e-01,  ..., -5.06076e+00, -2.24224e+00, -4.75443e+00],\n",
       "            [-9.41330e-01, -5.34662e-01, -4.14837e-01,  ..., -5.56409e+00, -3.18953e+00, -5.87509e+00],\n",
       "            [ 2.07958e-01, -1.38197e-01, -5.29104e-01,  ..., -5.64421e+00, -3.87367e+00, -5.93676e+00],\n",
       "            ...,\n",
       "            [-6.11877e-02, -2.73316e-01, -3.33043e-01,  ..., -5.66208e+00, -3.71443e+00, -5.76750e+00],\n",
       "            [ 4.75164e-04, -3.84414e-01, -4.92407e-01,  ..., -5.17263e+00, -3.94828e+00, -5.62433e+00],\n",
       "            [-1.42567e-01, -4.02105e-01, -4.46984e-01,  ..., -4.08961e+00, -2.46267e+00, -4.12437e+00]],\n",
       " \n",
       "           [[-6.04666e-01,  4.64475e-02, -8.54852e-01,  ..., -5.16834e+00, -2.90410e+00, -4.76797e+00],\n",
       "            [-1.03633e-01,  3.17011e-01, -7.80756e-01,  ..., -6.90040e+00, -4.21338e+00, -6.72841e+00],\n",
       "            [-6.92767e-03, -1.46995e-01, -6.59696e-01,  ..., -7.48095e+00, -5.01009e+00, -6.59553e+00],\n",
       "            ...,\n",
       "            [-6.93536e-01,  2.01018e-02, -8.19544e-01,  ..., -7.79615e+00, -5.86755e+00, -7.13492e+00],\n",
       "            [-7.88244e-02, -8.38478e-01, -4.81882e-01,  ..., -7.05231e+00, -4.71336e+00, -6.67126e+00],\n",
       "            [ 4.68860e-01,  2.09812e-01, -7.96226e-01,  ..., -5.13801e+00, -3.37647e+00, -5.54875e+00]],\n",
       " \n",
       "           [[-5.94911e-01,  6.96830e-02, -8.52138e-01,  ..., -5.20145e+00, -3.44268e+00, -4.72316e+00],\n",
       "            [-3.06270e-01,  1.84020e-01, -6.38567e-01,  ..., -7.32082e+00, -5.26573e+00, -6.39397e+00],\n",
       "            [ 5.29388e-01, -3.29561e-02, -4.29258e-01,  ..., -8.65974e+00, -5.67041e+00, -7.47241e+00],\n",
       "            ...,\n",
       "            [ 4.88498e-01, -2.09697e-01, -3.55003e-01,  ..., -8.75910e+00, -4.71639e+00, -8.16092e+00],\n",
       "            [-1.05642e+00,  1.72880e-01, -1.20675e-01,  ..., -8.60933e+00, -4.49691e+00, -8.12102e+00],\n",
       "            [-4.82291e-01, -1.86938e-01, -4.88558e-01,  ..., -5.97988e+00, -3.28196e+00, -5.95724e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-6.79332e-01,  1.50395e-01, -7.36383e-01,  ..., -5.36208e+00, -2.48901e+00, -4.57886e+00],\n",
       "            [ 5.04431e-01,  2.27552e-01, -4.84213e-01,  ..., -7.35466e+00, -4.43547e+00, -6.36767e+00],\n",
       "            [ 6.98082e-01, -8.30993e-01, -3.04782e-01,  ..., -7.94555e+00, -5.48215e+00, -7.08864e+00],\n",
       "            ...,\n",
       "            [-6.34688e-03,  3.38652e-02, -7.11823e-01,  ..., -7.79013e+00, -4.35295e+00, -6.86739e+00],\n",
       "            [-5.73866e-01, -9.60251e-02, -4.94753e-01,  ..., -7.25416e+00, -3.20093e+00, -6.90747e+00],\n",
       "            [ 2.31922e-01,  6.90083e-01, -6.38454e-01,  ..., -5.46944e+00, -2.21356e+00, -5.94147e+00]],\n",
       " \n",
       "           [[-4.73238e-01,  6.45130e-02, -7.54641e-01,  ..., -5.29853e+00, -3.10550e+00, -4.81652e+00],\n",
       "            [ 5.78812e-01, -5.91698e-01, -2.97716e-01,  ..., -6.89885e+00, -4.01894e+00, -6.14828e+00],\n",
       "            [ 9.19510e-01,  6.29096e-01, -2.64773e-01,  ..., -7.59772e+00, -5.04210e+00, -6.37962e+00],\n",
       "            ...,\n",
       "            [ 6.05812e-01, -6.64669e-01, -3.90089e-01,  ..., -6.57035e+00, -2.96208e+00, -6.50145e+00],\n",
       "            [-8.46333e-01, -1.84576e-02, -3.61114e-01,  ..., -6.88021e+00, -4.03103e+00, -7.20536e+00],\n",
       "            [ 5.10453e-01, -6.68899e-01, -8.08146e-01,  ..., -4.48622e+00, -2.25420e+00, -5.47461e+00]],\n",
       " \n",
       "           [[-2.19856e-02,  2.02054e-01, -5.59466e-01,  ..., -4.72369e+00, -4.29576e+00, -3.21947e+00],\n",
       "            [-4.70588e-01,  4.03704e-01, -2.35685e-01,  ..., -4.98132e+00, -4.39360e+00, -4.50213e+00],\n",
       "            [-3.74347e-01,  8.66703e-01, -4.55674e-02,  ..., -5.17538e+00, -4.31502e+00, -5.20729e+00],\n",
       "            ...,\n",
       "            [-2.05410e-01,  4.91215e-02, -2.94442e-01,  ..., -5.56093e+00, -3.83007e+00, -5.12155e+00],\n",
       "            [ 8.44422e-01,  2.89273e-01, -2.02695e-01,  ..., -4.73875e+00, -3.52470e+00, -4.60133e+00],\n",
       "            [-1.86284e-01,  1.61467e-01, -4.66870e-01,  ..., -4.19481e+00, -3.17017e+00, -3.39247e+00]]]]], device='cuda:0'))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_model(torch.rand((1,3,384,672)).cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d7821123",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mtrain: \u001b[0mScanning '../datasets/VisDrone/VisDrone2019-DET-train/labels.cache' images and labels... 6471 found, 0 missing, 0 empty, 4 corrupted: 100%|█| 6471/6471 \u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# For dataset\n",
    "imgsz = 640\n",
    "batch_size = 1\n",
    "gs = int(max(torch_model.stride))  # grid size (max stride)\n",
    "\n",
    "with open('data/hyp.scratch.tiny.yaml') as f:\n",
    "    hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps\n",
    "device='cuda'\n",
    "\n",
    "with open('./data/VisDrone.yaml') as f:\n",
    "    data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict\n",
    "    \n",
    "train_path = data_dict['train']\n",
    "test_path = data_dict['test']\n",
    "dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt=None,\n",
    "                                    hyp=hyp, augment=True, prefix=colorstr('train: '))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9f8e6c61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qat_model = torch_model.cuda()\n",
    "modified_state_dict = {}\n",
    "for key, val in qat_model.state_dict().items():\n",
    "\n",
    "    if key.startswith('module'):\n",
    "        modified_state_dict[key[7:]] = val\n",
    "    else:\n",
    "        modified_state_dict[key] = val\n",
    "\n",
    "qat_model.load_state_dict(modified_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6d1a8498",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[[[[-3.03100e-01,  8.63281e-03,  4.69530e-01,  ..., -5.78400e+00, -4.84398e+00, -3.23647e+00],\n",
       "            [-2.85701e-01, -8.69404e-02,  5.23454e-01,  ..., -7.12912e+00, -6.14127e+00, -4.98199e+00],\n",
       "            [-1.25635e+00, -5.33806e-02,  8.93788e-01,  ..., -6.50298e+00, -5.29050e+00, -3.29647e+00],\n",
       "            ...,\n",
       "            [-4.42532e-02,  2.11913e-02,  8.73779e-01,  ..., -7.55342e+00, -5.99426e+00, -5.13522e+00],\n",
       "            [-2.63270e-01, -1.73550e-02,  1.09284e+00,  ..., -7.57492e+00, -6.28809e+00, -5.00204e+00],\n",
       "            [-5.91200e-01, -2.54331e-02,  1.12599e+00,  ..., -6.36567e+00, -6.07072e+00, -4.09629e+00]],\n",
       " \n",
       "           [[-2.66897e-01, -8.99185e-02,  6.64318e-01,  ..., -5.55964e+00, -5.34481e+00, -3.88217e+00],\n",
       "            [-8.72903e-01, -3.89067e-01,  7.57212e-01,  ..., -6.28030e+00, -5.19728e+00, -4.38327e+00],\n",
       "            [-2.53817e-01, -3.04515e-01,  7.93453e-01,  ..., -6.85807e+00, -5.34068e+00, -4.34973e+00],\n",
       "            ...,\n",
       "            [-1.54081e-01, -8.41120e-01,  9.65253e-01,  ..., -7.27574e+00, -5.59024e+00, -4.43087e+00],\n",
       "            [-3.18531e-01, -1.03324e+00,  1.08304e+00,  ..., -7.43235e+00, -5.71875e+00, -4.53985e+00],\n",
       "            [-1.76947e-01, -1.30915e+00,  9.53132e-01,  ..., -5.83023e+00, -5.17962e+00, -2.75508e+00]],\n",
       " \n",
       "           [[-2.62891e-01, -2.20345e-01,  5.44550e-01,  ..., -6.65708e+00, -5.98856e+00, -5.15468e+00],\n",
       "            [-9.08957e-01,  9.37803e-02,  7.24782e-01,  ..., -6.68205e+00, -5.28038e+00, -4.69883e+00],\n",
       "            [-6.26276e-01,  1.00756e-02,  8.82797e-01,  ..., -6.30430e+00, -4.61381e+00, -3.96463e+00],\n",
       "            ...,\n",
       "            [-1.43158e-01, -5.35408e-01,  7.97977e-01,  ..., -6.91357e+00, -5.08876e+00, -4.02424e+00],\n",
       "            [-6.35275e-01, -5.93709e-01,  1.12970e+00,  ..., -6.46350e+00, -4.56115e+00, -2.95284e+00],\n",
       "            [ 2.29265e-01, -6.96157e-01,  6.35012e-01,  ..., -5.84645e+00, -4.78030e+00, -2.56007e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.00197e-01, -1.32284e-02,  4.93423e-01,  ..., -5.58563e+00, -4.39065e+00, -3.14156e+00],\n",
       "            [-6.93271e-01, -3.09198e-01,  9.73072e-01,  ..., -6.49055e+00, -4.62957e+00, -3.85542e+00],\n",
       "            [ 6.74749e-01, -1.01829e+00,  5.16636e-01,  ..., -7.71004e+00, -5.73644e+00, -5.29116e+00],\n",
       "            ...,\n",
       "            [-3.54477e-01, -5.62733e-01,  6.99257e-01,  ..., -6.38070e+00, -4.79868e+00, -3.59166e+00],\n",
       "            [-2.06064e-01,  8.62917e-01,  7.67860e-01,  ..., -6.64854e+00, -5.05854e+00, -3.98388e+00],\n",
       "            [-1.30458e-01,  5.25927e-01,  7.77488e-01,  ..., -5.86395e+00, -5.25755e+00, -2.96915e+00]],\n",
       " \n",
       "           [[-3.53356e-01,  3.52297e-01,  2.73167e-01,  ..., -5.96362e+00, -5.02445e+00, -4.05830e+00],\n",
       "            [-7.97927e-01,  8.38715e-01,  8.82923e-01,  ..., -5.33521e+00, -4.33488e+00, -3.47864e+00],\n",
       "            [-6.44233e-01,  5.00393e-01,  9.68877e-01,  ..., -6.23871e+00, -4.63739e+00, -3.96762e+00],\n",
       "            ...,\n",
       "            [-1.03166e+00,  9.18325e-01,  6.15022e-01,  ..., -6.85401e+00, -5.36719e+00, -4.04444e+00],\n",
       "            [-2.24038e-01,  3.83806e-01,  7.68949e-01,  ..., -7.10366e+00, -6.01340e+00, -3.84282e+00],\n",
       "            [ 3.29652e-01, -5.20593e-01,  5.85100e-01,  ..., -6.73315e+00, -6.67609e+00, -3.52326e+00]],\n",
       " \n",
       "           [[-5.30423e-01, -4.68569e-01,  6.26032e-01,  ..., -5.34457e+00, -5.52716e+00, -1.85406e+00],\n",
       "            [-8.21548e-01,  1.20427e-01,  7.32525e-01,  ..., -5.48355e+00, -4.91368e+00, -2.55407e+00],\n",
       "            [-7.52180e-01,  4.19492e-02,  9.17255e-01,  ..., -5.51866e+00, -4.66446e+00, -2.94585e+00],\n",
       "            ...,\n",
       "            [-7.72234e-01, -1.65652e-02,  8.38476e-01,  ..., -6.51602e+00, -5.68065e+00, -2.40028e+00],\n",
       "            [-4.90925e-01, -1.55114e-01,  6.80892e-01,  ..., -7.72148e+00, -6.42526e+00, -3.93323e+00],\n",
       "            [-3.67258e-01, -3.24888e-01,  8.50291e-01,  ..., -7.20462e+00, -6.47303e+00, -3.24326e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 6.16520e-01,  5.72757e-02,  4.32767e-01,  ..., -5.70545e+00, -5.05219e+00, -2.35469e+00],\n",
       "            [ 4.47354e-01,  3.75444e-02,  3.13230e-01,  ..., -7.90824e+00, -6.63382e+00, -3.88290e+00],\n",
       "            [-8.14523e-02,  8.50245e-02,  8.83782e-01,  ..., -6.94720e+00, -5.53689e+00, -2.44052e+00],\n",
       "            ...,\n",
       "            [ 6.69736e-01,  3.34037e-01,  6.70043e-01,  ..., -8.35007e+00, -6.09024e+00, -4.16372e+00],\n",
       "            [ 9.30027e-01,  2.34636e-01,  7.94149e-01,  ..., -8.28264e+00, -6.60844e+00, -4.35881e+00],\n",
       "            [-2.79023e-02,  1.38706e-01,  6.45481e-01,  ..., -6.22390e+00, -6.27701e+00, -3.41032e+00]],\n",
       " \n",
       "           [[-2.67070e-02,  4.90507e-01,  5.76663e-01,  ..., -6.09171e+00, -5.60313e+00, -3.18076e+00],\n",
       "            [ 3.73900e-01, -2.31863e-01,  7.19698e-01,  ..., -6.98591e+00, -5.69302e+00, -3.87974e+00],\n",
       "            [ 6.74087e-01, -7.74580e-02,  8.87341e-01,  ..., -7.37250e+00, -5.56387e+00, -4.19802e+00],\n",
       "            ...,\n",
       "            [ 6.49868e-01, -5.38131e-01,  1.06593e+00,  ..., -7.88815e+00, -5.84572e+00, -4.28777e+00],\n",
       "            [ 5.42189e-01, -8.29755e-01,  1.15603e+00,  ..., -7.81300e+00, -5.94988e+00, -4.27963e+00],\n",
       "            [ 2.07421e-01, -8.46198e-01,  6.28929e-01,  ..., -5.46535e+00, -4.93973e+00, -2.89219e+00]],\n",
       " \n",
       "           [[-2.88333e-01,  1.86820e-01,  3.82669e-01,  ..., -7.03428e+00, -6.03040e+00, -4.69337e+00],\n",
       "            [ 2.15616e-01,  1.86878e-01,  6.80273e-01,  ..., -7.38145e+00, -5.51767e+00, -4.54652e+00],\n",
       "            [ 4.71860e-01,  1.91669e-01,  9.62887e-01,  ..., -6.73712e+00, -4.91350e+00, -3.86953e+00],\n",
       "            ...,\n",
       "            [ 1.99039e-01, -3.16087e-01,  9.15880e-01,  ..., -7.42448e+00, -5.17841e+00, -4.17583e+00],\n",
       "            [ 5.41890e-01, -2.80790e-01,  1.46705e+00,  ..., -6.83209e+00, -4.62026e+00, -2.93598e+00],\n",
       "            [ 4.22628e-01, -3.16413e-01,  3.85097e-01,  ..., -5.71415e+00, -4.29243e+00, -2.63296e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.18112e-02,  3.23498e-01,  4.91021e-01,  ..., -5.65138e+00, -4.21367e+00, -3.02419e+00],\n",
       "            [ 5.58424e-01, -3.14050e-01,  1.07181e+00,  ..., -7.03814e+00, -4.76924e+00, -3.97036e+00],\n",
       "            [ 9.87419e-01, -8.07472e-01,  4.99775e-01,  ..., -8.50688e+00, -5.69336e+00, -4.85608e+00],\n",
       "            ...,\n",
       "            [ 4.43258e-01, -2.46062e-01,  7.79379e-01,  ..., -7.07260e+00, -5.04470e+00, -3.72139e+00],\n",
       "            [ 1.30290e+00,  1.06034e+00,  8.18402e-01,  ..., -7.44719e+00, -5.15211e+00, -3.69495e+00],\n",
       "            [ 2.20410e-01,  9.40310e-01,  4.66210e-01,  ..., -6.05194e+00, -5.08639e+00, -2.63528e+00]],\n",
       " \n",
       "           [[-2.25932e-01,  9.27697e-01,  2.09437e-01,  ..., -6.22083e+00, -4.94327e+00, -3.59076e+00],\n",
       "            [ 8.89987e-01,  1.03580e+00,  9.10666e-01,  ..., -5.68750e+00, -4.38910e+00, -3.15930e+00],\n",
       "            [ 3.33675e-01,  9.00482e-01,  8.88589e-01,  ..., -7.08539e+00, -4.46203e+00, -3.75354e+00],\n",
       "            ...,\n",
       "            [ 1.41082e-02,  1.31933e+00,  6.87995e-01,  ..., -7.98316e+00, -5.37582e+00, -3.73723e+00],\n",
       "            [ 8.20680e-01,  7.64491e-01,  7.11978e-01,  ..., -8.14499e+00, -6.05239e+00, -3.60493e+00],\n",
       "            [ 2.67164e-01,  6.99444e-02,  2.91078e-01,  ..., -7.67406e+00, -6.57218e+00, -2.96643e+00]],\n",
       " \n",
       "           [[-7.33899e-02, -3.86755e-01,  5.17197e-01,  ..., -5.35072e+00, -5.42200e+00, -1.85649e+00],\n",
       "            [ 3.54604e-01,  4.35998e-02,  5.91259e-01,  ..., -5.82330e+00, -4.89486e+00, -2.34104e+00],\n",
       "            [ 9.35226e-02,  4.41763e-02,  6.95103e-01,  ..., -5.84469e+00, -4.45368e+00, -2.72060e+00],\n",
       "            ...,\n",
       "            [ 7.53881e-01, -1.27216e-01,  6.07271e-01,  ..., -7.16080e+00, -5.62996e+00, -2.44890e+00],\n",
       "            [ 2.78807e-01, -2.63145e-01,  4.52209e-01,  ..., -8.85233e+00, -6.67774e+00, -3.76300e+00],\n",
       "            [ 1.51018e-02, -3.81637e-01,  4.88340e-01,  ..., -7.50891e+00, -6.60366e+00, -3.05772e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 1.82030e-01,  2.25813e-01, -7.91429e-02,  ..., -4.07155e+00, -4.56971e+00, -1.44660e+00],\n",
       "            [ 4.95896e-02, -8.70816e-03,  2.08643e-02,  ..., -6.26593e+00, -5.77036e+00, -2.89991e+00],\n",
       "            [-6.97892e-01,  1.14114e-01,  5.05295e-01,  ..., -5.19352e+00, -4.60066e+00, -2.38319e+00],\n",
       "            ...,\n",
       "            [ 2.31602e-01,  3.73273e-01,  2.54769e-01,  ..., -6.94100e+00, -5.26575e+00, -3.80537e+00],\n",
       "            [ 2.12684e-01,  2.65365e-01,  2.50981e-01,  ..., -6.62868e+00, -5.71816e+00, -3.41593e+00],\n",
       "            [-4.37851e-01,  2.21390e-01,  9.07250e-02,  ..., -4.74590e+00, -5.95042e+00, -2.57815e+00]],\n",
       " \n",
       "           [[-1.19381e-01,  1.08700e+00, -1.24801e-01,  ..., -5.25203e+00, -5.06442e+00, -2.82994e+00],\n",
       "            [-4.38783e-01,  1.80800e-01,  3.33903e-01,  ..., -6.12119e+00, -4.58884e+00, -3.71583e+00],\n",
       "            [-1.67069e-01,  2.39963e-01,  4.58121e-01,  ..., -6.36459e+00, -4.66335e+00, -4.38244e+00],\n",
       "            ...,\n",
       "            [ 1.14615e-01, -2.17446e-01,  5.13178e-01,  ..., -7.29466e+00, -4.92153e+00, -4.55040e+00],\n",
       "            [-1.77249e-01, -5.65522e-01,  5.73526e-01,  ..., -6.86801e+00, -5.23540e+00, -4.05586e+00],\n",
       "            [-2.45418e-01, -5.19198e-01,  2.11796e-01,  ..., -4.53896e+00, -4.67890e+00, -2.63638e+00]],\n",
       " \n",
       "           [[-2.39634e-01,  5.76554e-01, -2.96609e-01,  ..., -6.15385e+00, -5.72424e+00, -4.10540e+00],\n",
       "            [-3.93744e-01,  3.93737e-01,  1.83808e-01,  ..., -6.31202e+00, -4.54270e+00, -4.67336e+00],\n",
       "            [-2.37121e-01,  5.55657e-01,  3.89002e-01,  ..., -5.85164e+00, -3.99503e+00, -4.30515e+00],\n",
       "            ...,\n",
       "            [-1.71052e-01, -4.01721e-02,  3.67268e-01,  ..., -6.93998e+00, -4.53429e+00, -4.40414e+00],\n",
       "            [-3.05466e-01,  1.53358e-01,  7.83602e-01,  ..., -6.38388e+00, -3.78624e+00, -3.62555e+00],\n",
       "            [-6.07968e-03,  1.27352e-01,  1.03877e-02,  ..., -4.82875e+00, -4.11768e+00, -2.79967e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.31811e-01,  6.70909e-01, -1.29305e-01,  ..., -4.60114e+00, -3.95479e+00, -3.48282e+00],\n",
       "            [-3.08696e-01, -3.88622e-02,  5.04035e-01,  ..., -6.12514e+00, -3.84618e+00, -5.03137e+00],\n",
       "            [ 7.91749e-01, -4.94465e-01,  4.84507e-02,  ..., -7.97274e+00, -5.04478e+00, -5.51156e+00],\n",
       "            ...,\n",
       "            [-3.63842e-02,  3.80076e-02,  2.50971e-01,  ..., -6.75860e+00, -4.30704e+00, -4.62482e+00],\n",
       "            [ 5.03957e-01,  1.22194e+00,  2.67537e-01,  ..., -6.78777e+00, -4.37312e+00, -4.26864e+00],\n",
       "            [-5.09045e-02,  1.24992e+00, -9.15035e-02,  ..., -5.32273e+00, -4.75832e+00, -2.84092e+00]],\n",
       " \n",
       "           [[-2.02004e-01,  1.43508e+00, -3.83860e-01,  ..., -5.35371e+00, -4.77667e+00, -3.47556e+00],\n",
       "            [ 6.01264e-02,  1.32606e+00,  3.91707e-01,  ..., -4.77531e+00, -3.46772e+00, -3.54303e+00],\n",
       "            [-3.81203e-01,  1.39570e+00,  3.36565e-01,  ..., -6.50745e+00, -3.61800e+00, -5.21207e+00],\n",
       "            ...,\n",
       "            [-5.61658e-01,  1.70941e+00,  1.47525e-01,  ..., -7.32941e+00, -4.56110e+00, -4.90526e+00],\n",
       "            [ 1.46293e-01,  1.29883e+00,  1.02052e-01,  ..., -7.62343e+00, -5.28380e+00, -4.14810e+00],\n",
       "            [ 2.08946e-01,  7.87498e-01, -2.63168e-01,  ..., -7.29659e+00, -6.17043e+00, -2.72320e+00]],\n",
       " \n",
       "           [[-7.91299e-02, -1.68466e-01, -1.08797e-01,  ..., -4.84394e+00, -5.39026e+00, -1.72691e+00],\n",
       "            [-4.27673e-01,  6.29765e-02,  1.36913e-01,  ..., -5.18284e+00, -4.39931e+00, -2.63076e+00],\n",
       "            [-4.55406e-01,  2.02064e-01,  1.92698e-01,  ..., -5.32496e+00, -4.13880e+00, -2.86681e+00],\n",
       "            ...,\n",
       "            [-1.10213e-01, -2.16232e-03,  4.20572e-02,  ..., -6.46107e+00, -5.08961e+00, -2.71209e+00],\n",
       "            [-1.91179e-01, -2.64061e-01, -1.93298e-01,  ..., -8.07112e+00, -6.19137e+00, -3.77983e+00],\n",
       "            [-1.88472e-01, -3.46642e-01, -1.05308e-01,  ..., -6.77070e+00, -6.43978e+00, -2.96333e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-1.69511e-01, -9.88329e-02,  3.55242e-01,  ..., -6.15250e+00, -3.18090e+00, -5.78410e+00],\n",
       "            [-1.86005e-01, -3.29535e-01,  8.56919e-01,  ..., -7.31080e+00, -3.50155e+00, -5.61154e+00],\n",
       "            [-2.60967e-01, -3.84464e-01,  5.64130e-01,  ..., -7.86901e+00, -4.19429e+00, -5.66973e+00],\n",
       "            ...,\n",
       "            [ 2.65221e-01, -5.08938e-01,  7.09928e-01,  ..., -7.11694e+00, -3.86802e+00, -5.32968e+00],\n",
       "            [ 3.23323e-01, -2.29024e-01,  1.00549e+00,  ..., -7.60336e+00, -3.08431e+00, -5.70124e+00],\n",
       "            [ 2.32796e-01, -2.15021e-01,  4.25407e-01,  ..., -6.77449e+00, -2.96697e+00, -6.08618e+00]],\n",
       " \n",
       "           [[-5.38787e-01, -8.74747e-02, -1.01718e-01,  ..., -8.05604e+00, -4.75537e+00, -6.44586e+00],\n",
       "            [-6.30918e-02,  7.16263e-01,  4.58401e-01,  ..., -8.95120e+00, -4.90272e+00, -5.93605e+00],\n",
       "            [ 7.23005e-02,  8.22366e-01,  7.61639e-01,  ..., -8.50946e+00, -4.05964e+00, -5.39826e+00],\n",
       "            ...,\n",
       "            [ 3.75889e-02,  9.67501e-01,  7.60490e-01,  ..., -8.05856e+00, -4.44826e+00, -5.50375e+00],\n",
       "            [ 3.64893e-01, -4.23542e-01,  6.76850e-01,  ..., -8.24352e+00, -4.49698e+00, -5.26478e+00],\n",
       "            [ 4.87884e-01, -2.30123e-01,  9.73149e-02,  ..., -7.52347e+00, -3.47846e+00, -5.73429e+00]],\n",
       " \n",
       "           [[-4.75315e-01, -2.13090e-01, -3.77130e-02,  ..., -8.77870e+00, -4.20247e+00, -6.38813e+00],\n",
       "            [ 2.70960e-01,  1.96634e-01,  4.47844e-01,  ..., -1.01681e+01, -5.17062e+00, -6.10575e+00],\n",
       "            [-8.49679e-01, -2.56130e-01,  5.08468e-01,  ..., -9.31702e+00, -4.15503e+00, -6.02649e+00],\n",
       "            ...,\n",
       "            [-1.28404e-01, -4.34671e-01,  4.83096e-01,  ..., -8.90227e+00, -4.62575e+00, -6.14242e+00],\n",
       "            [-3.04147e-01,  4.30953e-01,  4.64747e-01,  ..., -9.40526e+00, -5.05130e+00, -6.19835e+00],\n",
       "            [ 2.64544e-01,  6.41903e-01,  2.59030e-01,  ..., -8.28722e+00, -3.41948e+00, -6.10639e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.94465e-01,  5.88907e-01,  7.56478e-03,  ..., -8.84885e+00, -2.73657e+00, -5.92865e+00],\n",
       "            [-1.67409e-02,  1.06236e+00,  4.65413e-01,  ..., -1.07785e+01, -4.53844e+00, -6.87924e+00],\n",
       "            [-6.75614e-01,  6.18978e-01,  3.23176e-01,  ..., -1.06561e+01, -6.05016e+00, -6.75308e+00],\n",
       "            ...,\n",
       "            [ 4.97025e-01, -2.93638e-01,  3.63765e-01,  ..., -9.34835e+00, -5.27364e+00, -5.70560e+00],\n",
       "            [ 4.05452e-01, -2.19740e-01,  3.66855e-01,  ..., -9.55378e+00, -4.76291e+00, -5.68403e+00],\n",
       "            [ 4.03772e-01,  4.65601e-01,  2.28951e-01,  ..., -7.41472e+00, -3.11945e+00, -5.26402e+00]],\n",
       " \n",
       "           [[-5.22111e-01,  9.09582e-01,  6.59965e-02,  ..., -7.92319e+00, -3.03828e+00, -5.72224e+00],\n",
       "            [-1.42340e-01, -7.04691e-01,  2.86567e-01,  ..., -1.10577e+01, -4.98208e+00, -7.09670e+00],\n",
       "            [-7.57259e-01, -7.79212e-01,  4.21648e-01,  ..., -1.00612e+01, -4.29518e+00, -6.28418e+00],\n",
       "            ...,\n",
       "            [ 4.71113e-01,  3.11237e-01,  4.53047e-01,  ..., -9.20426e+00, -4.41395e+00, -5.45069e+00],\n",
       "            [ 3.85780e-01, -1.08814e-01,  5.46264e-01,  ..., -9.12561e+00, -4.50447e+00, -5.62165e+00],\n",
       "            [ 2.38948e-01,  1.07220e+00,  2.71876e-01,  ..., -7.01586e+00, -2.86467e+00, -5.66632e+00]],\n",
       " \n",
       "           [[-4.14790e-01,  1.64506e-01,  5.11400e-02,  ..., -6.90866e+00, -4.40637e+00, -6.62071e+00],\n",
       "            [-2.36412e-01,  4.16039e-01,  6.18393e-01,  ..., -7.93852e+00, -3.31848e+00, -6.11092e+00],\n",
       "            [ 1.02347e-01,  5.75208e-01,  7.39945e-01,  ..., -7.94197e+00, -2.85257e+00, -6.08219e+00],\n",
       "            ...,\n",
       "            [-3.61998e-01,  2.42601e-01,  3.98935e-01,  ..., -7.72134e+00, -4.72471e+00, -6.07420e+00],\n",
       "            [ 2.70128e-01,  2.31134e-01,  2.18454e-01,  ..., -7.73734e+00, -4.80494e+00, -6.19664e+00],\n",
       "            [ 3.07683e-01,  1.61980e-01,  1.88153e-01,  ..., -7.21854e+00, -3.81952e+00, -6.87605e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 4.60340e-02, -2.48658e-01, -5.09170e-02,  ..., -5.39878e+00, -3.14102e+00, -3.40013e+00],\n",
       "            [-1.08894e-01, -4.47280e-01,  4.75165e-01,  ..., -5.08831e+00, -2.98446e+00, -4.43042e+00],\n",
       "            [-2.43253e-01, -4.81186e-01,  2.91634e-01,  ..., -5.17487e+00, -3.01529e+00, -4.13868e+00],\n",
       "            ...,\n",
       "            [ 3.12353e-01, -5.93437e-01,  3.13633e-01,  ..., -4.99120e+00, -2.95348e+00, -4.47100e+00],\n",
       "            [ 3.39066e-01, -3.37493e-01,  3.99095e-01,  ..., -4.93968e+00, -2.48536e+00, -5.12931e+00],\n",
       "            [ 1.33681e-01, -3.10037e-01, -1.29945e-01,  ..., -5.60124e+00, -2.54539e+00, -3.54628e+00]],\n",
       " \n",
       "           [[-4.35396e-01, -1.91991e-01, -5.01767e-01,  ..., -5.65798e+00, -3.93936e+00, -4.33557e+00],\n",
       "            [-1.57263e-01,  5.78252e-01,  1.16108e-01,  ..., -4.92600e+00, -3.87157e+00, -6.24179e+00],\n",
       "            [ 3.40059e-02,  6.30655e-01,  2.37878e-01,  ..., -4.52480e+00, -2.91266e+00, -5.75682e+00],\n",
       "            ...,\n",
       "            [ 1.23362e-02,  8.24793e-01,  2.13815e-01,  ..., -4.54215e+00, -3.03056e+00, -6.14048e+00],\n",
       "            [ 2.98012e-01, -4.89100e-01,  2.34317e-01,  ..., -4.64938e+00, -3.53929e+00, -6.02272e+00],\n",
       "            [ 4.49337e-01, -3.42999e-01, -3.36515e-01,  ..., -5.22397e+00, -3.05673e+00, -4.90661e+00]],\n",
       " \n",
       "           [[-4.22399e-01, -3.16915e-01, -4.85159e-01,  ..., -5.57930e+00, -3.19008e+00, -4.85416e+00],\n",
       "            [ 2.34632e-01, -1.21942e-01, -7.54913e-02,  ..., -5.21974e+00, -3.22161e+00, -6.06632e+00],\n",
       "            [-8.95706e-01, -3.43836e-01, -3.26290e-02,  ..., -5.05395e+00, -2.82713e+00, -4.76870e+00],\n",
       "            ...,\n",
       "            [-1.15983e-01, -4.74510e-01,  4.16866e-02,  ..., -5.17795e+00, -3.10899e+00, -5.57854e+00],\n",
       "            [-2.89974e-01,  3.55385e-01,  4.29281e-02,  ..., -5.10516e+00, -3.69961e+00, -5.70344e+00],\n",
       "            [ 1.97705e-01,  5.37408e-01, -1.98901e-01,  ..., -5.16686e+00, -2.81325e+00, -5.94414e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.82898e-01,  4.40198e-01, -4.52679e-01,  ..., -5.18853e+00, -2.40596e+00, -6.25791e+00],\n",
       "            [-4.47472e-02,  9.66129e-01, -1.38383e-01,  ..., -5.45870e+00, -2.64232e+00, -6.89449e+00],\n",
       "            [-7.23655e-01,  5.45895e-01, -1.52343e-01,  ..., -5.49500e+00, -3.78118e+00, -6.00746e+00],\n",
       "            ...,\n",
       "            [ 5.03978e-01, -3.62230e-01, -7.92157e-02,  ..., -4.86830e+00, -2.94765e+00, -7.54869e+00],\n",
       "            [ 4.48036e-01, -2.87288e-01, -5.11775e-02,  ..., -4.88694e+00, -2.91998e+00, -6.24941e+00],\n",
       "            [ 3.72060e-01,  2.62569e-01, -1.69812e-01,  ..., -5.04952e+00, -2.89339e+00, -5.83244e+00]],\n",
       " \n",
       "           [[-4.72819e-01,  7.06030e-01, -3.98395e-01,  ..., -5.14794e+00, -2.85324e+00, -4.83813e+00],\n",
       "            [-1.47954e-01, -7.38940e-01, -2.45393e-01,  ..., -5.88780e+00, -3.56486e+00, -4.84217e+00],\n",
       "            [-8.65009e-01, -8.07394e-01, -1.09192e-01,  ..., -5.50839e+00, -3.21432e+00, -5.44577e+00],\n",
       "            ...,\n",
       "            [ 4.88548e-01,  1.85238e-01,  3.18831e-02,  ..., -4.58995e+00, -3.03780e+00, -6.63203e+00],\n",
       "            [ 4.32517e-01, -1.92793e-01,  1.32951e-01,  ..., -4.84208e+00, -3.26549e+00, -6.03001e+00],\n",
       "            [ 1.71008e-01,  7.86868e-01, -2.64155e-01,  ..., -5.25917e+00, -2.71853e+00, -5.11110e+00]],\n",
       " \n",
       "           [[-2.43139e-01,  1.74410e-01, -3.16423e-01,  ..., -5.85242e+00, -3.82428e+00, -2.37595e+00],\n",
       "            [-2.75112e-01,  3.46613e-01,  2.92261e-01,  ..., -5.24569e+00, -2.85449e+00, -5.35841e+00],\n",
       "            [ 1.21822e-01,  4.53289e-01,  2.58347e-01,  ..., -5.31433e+00, -2.45493e+00, -5.39664e+00],\n",
       "            ...,\n",
       "            [-3.24877e-01,  1.87164e-01,  1.25239e-01,  ..., -5.20416e+00, -3.78744e+00, -4.27814e+00],\n",
       "            [ 2.79416e-01,  1.49601e-01,  1.55293e-02,  ..., -5.34968e+00, -3.59696e+00, -3.77344e+00],\n",
       "            [ 2.52096e-01,  1.32363e-01, -3.92236e-01,  ..., -6.10994e+00, -2.93654e+00, -2.24296e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 5.20700e-02, -3.50107e-02,  1.68917e-01,  ..., -4.80725e+00, -2.30755e+00, -3.33117e+00],\n",
       "            [-1.06511e-01, -2.77792e-01,  8.47775e-01,  ..., -5.70321e+00, -2.75171e+00, -4.27188e+00],\n",
       "            [-2.31578e-01, -3.47197e-01,  6.33189e-01,  ..., -6.05542e+00, -3.24130e+00, -4.50545e+00],\n",
       "            ...,\n",
       "            [ 3.27080e-01, -4.64659e-01,  6.33365e-01,  ..., -5.65477e+00, -3.01356e+00, -4.62108e+00],\n",
       "            [ 3.36970e-01, -1.95308e-01,  7.30379e-01,  ..., -6.15697e+00, -2.27118e+00, -5.23196e+00],\n",
       "            [ 1.27297e-01, -1.81192e-01,  9.50401e-02,  ..., -5.44493e+00, -1.98716e+00, -3.75096e+00]],\n",
       " \n",
       "           [[-4.22603e-01, -9.02727e-02, -3.64763e-01,  ..., -5.93028e+00, -3.73123e+00, -3.54420e+00],\n",
       "            [-1.39261e-01,  6.71880e-01,  4.05562e-01,  ..., -6.91126e+00, -4.22734e+00, -5.09347e+00],\n",
       "            [ 5.07125e-02,  7.03098e-01,  5.71018e-01,  ..., -6.90254e+00, -3.04841e+00, -6.20806e+00],\n",
       "            ...,\n",
       "            [ 3.79598e-02,  8.98861e-01,  5.02058e-01,  ..., -6.70498e+00, -3.45950e+00, -6.85881e+00],\n",
       "            [ 3.12379e-01, -4.57412e-01,  5.63850e-01,  ..., -6.71533e+00, -3.64786e+00, -5.62389e+00],\n",
       "            [ 4.41776e-01, -2.88735e-01, -1.02295e-01,  ..., -5.93743e+00, -2.41245e+00, -5.02964e+00]],\n",
       " \n",
       "           [[-4.05398e-01, -2.18175e-01, -3.50180e-01,  ..., -6.81073e+00, -3.02251e+00, -4.91261e+00],\n",
       "            [ 2.61619e-01, -3.62294e-02,  1.24789e-01,  ..., -7.76357e+00, -3.71402e+00, -6.79779e+00],\n",
       "            [-8.52659e-01, -2.87985e-01,  1.83494e-01,  ..., -7.78280e+00, -3.03712e+00, -7.77009e+00],\n",
       "            ...,\n",
       "            [-9.40306e-02, -4.73399e-01,  3.10728e-01,  ..., -7.63430e+00, -3.57284e+00, -6.26528e+00],\n",
       "            [-2.61514e-01,  4.29550e-01,  2.99845e-01,  ..., -7.48772e+00, -4.07184e+00, -5.85957e+00],\n",
       "            [ 2.02122e-01,  6.75479e-01,  2.73711e-02,  ..., -6.99887e+00, -2.74843e+00, -5.29740e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.59081e-01,  4.94458e-01, -3.07656e-01,  ..., -7.59213e+00, -1.73208e+00, -6.27915e+00],\n",
       "            [-3.03917e-02,  1.00759e+00,  4.57241e-02,  ..., -8.97037e+00, -3.01020e+00, -8.18163e+00],\n",
       "            [-6.94991e-01,  5.85407e-01,  2.29770e-02,  ..., -8.41991e+00, -4.76543e+00, -6.82855e+00],\n",
       "            ...,\n",
       "            [ 5.20040e-01, -3.52992e-01,  1.68358e-01,  ..., -8.30038e+00, -4.21375e+00, -7.89308e+00],\n",
       "            [ 4.68271e-01, -2.82495e-01,  2.16196e-01,  ..., -8.03077e+00, -3.51887e+00, -7.58017e+00],\n",
       "            [ 3.79012e-01,  3.24106e-01,  1.23102e-01,  ..., -6.50820e+00, -2.08834e+00, -5.59261e+00]],\n",
       " \n",
       "           [[-4.46653e-01,  7.80273e-01, -2.50402e-01,  ..., -6.58497e+00, -2.01877e+00, -4.92876e+00],\n",
       "            [-1.27425e-01, -7.31323e-01, -6.44420e-02,  ..., -8.41085e+00, -3.80440e+00, -5.40145e+00],\n",
       "            [-8.34172e-01, -7.67454e-01,  1.12846e-01,  ..., -8.07201e+00, -3.21222e+00, -6.61892e+00],\n",
       "            ...,\n",
       "            [ 5.05418e-01,  2.40790e-01,  2.83596e-01,  ..., -7.92356e+00, -3.62387e+00, -7.26723e+00],\n",
       "            [ 4.44299e-01, -1.37542e-01,  4.42543e-01,  ..., -7.71209e+00, -3.49923e+00, -6.42049e+00],\n",
       "            [ 1.69469e-01,  9.33457e-01, -2.71648e-02,  ..., -6.23366e+00, -1.77176e+00, -5.05233e+00]],\n",
       " \n",
       "           [[-2.36501e-01,  1.50690e-01, -1.57203e-01,  ..., -5.06140e+00, -3.54006e+00, -1.71976e+00],\n",
       "            [-2.60100e-01,  3.02920e-01,  6.59133e-01,  ..., -6.48333e+00, -2.46088e+00, -4.21189e+00],\n",
       "            [ 1.38870e-01,  4.04128e-01,  5.86963e-01,  ..., -6.70114e+00, -1.74171e+00, -5.15136e+00],\n",
       "            ...,\n",
       "            [-3.09486e-01,  2.08616e-01,  4.27625e-01,  ..., -6.20597e+00, -3.93244e+00, -3.75826e+00],\n",
       "            [ 2.94035e-01,  1.32310e-01,  3.36530e-01,  ..., -5.80024e+00, -3.69523e+00, -3.31778e+00],\n",
       "            [ 2.44491e-01,  9.93916e-02, -2.07127e-01,  ..., -5.58546e+00, -2.76047e+00, -2.56698e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-2.06286e-01, -2.24522e-02, -3.88936e-01,  ..., -5.40068e+00, -3.90216e+00, -2.80304e+00],\n",
       "            [-1.16138e+00, -3.00273e-01, -3.67932e-02,  ..., -5.20965e+00, -4.25457e+00, -2.76036e+00],\n",
       "            [ 2.60514e-01,  1.17335e-01, -3.86127e-01,  ..., -5.78850e+00, -5.96142e+00, -1.92686e+00],\n",
       "            ...,\n",
       "            [ 1.39701e-01, -4.29670e-01, -2.39987e-01,  ..., -5.63806e+00, -5.56067e+00, -2.52426e+00],\n",
       "            [ 1.10610e-01, -7.85271e-01,  1.35264e-01,  ..., -5.48537e+00, -4.41196e+00, -3.93554e+00],\n",
       "            [ 2.98662e-01,  1.37742e-01, -2.33659e-01,  ..., -5.02663e+00, -2.84248e+00, -3.58959e+00]],\n",
       " \n",
       "           [[-4.62304e-01,  9.64274e-02, -7.04505e-01,  ..., -5.55570e+00, -5.00903e+00, -2.64004e+00],\n",
       "            [ 1.99866e-02, -3.27491e-01, -7.88338e-01,  ..., -6.02747e+00, -6.90219e+00, -3.02210e+00],\n",
       "            [-1.33076e-01, -8.69382e-01, -6.56561e-01,  ..., -6.08052e+00, -7.42014e+00, -3.36545e+00],\n",
       "            ...,\n",
       "            [ 2.18521e-01, -3.69865e-01, -5.26013e-01,  ..., -5.82541e+00, -6.78777e+00, -3.11613e+00],\n",
       "            [-9.54183e-01,  3.07158e-01, -5.39910e-01,  ..., -6.09464e+00, -6.95116e+00, -3.66152e+00],\n",
       "            [ 2.64339e-01,  2.00062e-01, -5.48192e-01,  ..., -5.36752e+00, -4.19748e+00, -3.03996e+00]],\n",
       " \n",
       "           [[-4.07466e-01,  2.43179e-01, -6.71089e-01,  ..., -5.44940e+00, -4.82803e+00, -1.18512e+00],\n",
       "            [-2.92394e-01,  2.78784e-01, -6.07225e-01,  ..., -6.18596e+00, -6.41510e+00, -3.01794e+00],\n",
       "            [-8.15060e-01,  1.79654e-01, -5.23872e-01,  ..., -6.59231e+00, -8.71394e+00, -3.40902e+00],\n",
       "            ...,\n",
       "            [-1.08729e+00,  9.68620e-02, -4.46952e-01,  ..., -6.36144e+00, -8.10883e+00, -3.14576e+00],\n",
       "            [ 4.75848e-01,  3.07574e-01, -4.43488e-01,  ..., -6.01078e+00, -6.79226e+00, -2.93710e+00],\n",
       "            [-2.69466e-01, -1.51617e-02, -4.51249e-01,  ..., -5.30945e+00, -5.14314e+00, -2.75465e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-8.93280e-01,  2.32601e-01, -9.07208e-01,  ..., -5.75673e+00, -3.79566e+00, -3.76293e+00],\n",
       "            [ 8.47292e-01, -2.04671e-01, -6.17289e-01,  ..., -6.37256e+00, -6.67231e+00, -3.46036e+00],\n",
       "            [ 6.14026e-01, -8.84967e-01, -3.39178e-01,  ..., -6.73208e+00, -8.10772e+00, -2.92501e+00],\n",
       "            ...,\n",
       "            [-1.03936e+00,  6.12278e-02, -5.06281e-01,  ..., -6.55179e+00, -7.64769e+00, -2.84085e+00],\n",
       "            [-9.75497e-01, -6.10315e-01, -3.56168e-01,  ..., -5.91075e+00, -6.87986e+00, -4.00195e+00],\n",
       "            [ 7.16478e-01,  6.45377e-01, -4.83058e-01,  ..., -5.87205e+00, -3.85117e+00, -4.09886e+00]],\n",
       " \n",
       "           [[-1.82059e-02,  6.19824e-01, -7.64264e-01,  ..., -5.78285e+00, -4.20899e+00, -3.27618e+00],\n",
       "            [ 3.34287e-02, -6.02081e-02, -6.20445e-01,  ..., -5.68376e+00, -4.87916e+00, -3.57272e+00],\n",
       "            [ 2.71763e-01, -1.02697e+00, -5.44701e-01,  ..., -6.11347e+00, -6.68551e+00, -3.33918e+00],\n",
       "            ...,\n",
       "            [-4.97148e-01, -8.22626e-01, -3.40411e-02,  ..., -5.78445e+00, -6.29254e+00, -4.23990e+00],\n",
       "            [-2.64134e-02, -7.14760e-01,  9.04844e-02,  ..., -5.49071e+00, -6.04896e+00, -4.50875e+00],\n",
       "            [ 3.27885e-01,  6.39377e-01, -2.21929e-01,  ..., -5.43953e+00, -4.27676e+00, -3.98443e+00]],\n",
       " \n",
       "           [[-6.67287e-02,  3.92121e-01, -3.17003e-01,  ..., -4.83895e+00, -3.65465e+00, -2.52329e+00],\n",
       "            [ 6.23237e-03,  5.09604e-01, -2.46811e-01,  ..., -5.06280e+00, -4.57021e+00, -2.61070e+00],\n",
       "            [-3.97327e-01,  1.81283e-01, -4.64341e-01,  ..., -5.27617e+00, -6.03123e+00, -2.65359e+00],\n",
       "            ...,\n",
       "            [-8.18302e-01,  6.30059e-01, -7.05839e-02,  ..., -4.95000e+00, -4.43468e+00, -2.20590e+00],\n",
       "            [-5.75280e-01,  3.71119e-01, -5.86298e-02,  ..., -4.87087e+00, -4.68758e+00, -2.52970e+00],\n",
       "            [-7.53641e-02,  3.45784e-01, -3.88866e-02,  ..., -4.81760e+00, -4.04312e+00, -2.51427e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.16557e-01, -7.43256e-03, -3.21113e-01,  ..., -4.52383e+00, -4.73988e+00, -6.25103e+00],\n",
       "            [-1.14620e+00, -2.67338e-01,  4.29849e-02,  ..., -5.62861e+00, -5.27354e+00, -7.04421e+00],\n",
       "            [ 2.72970e-01,  1.57926e-01, -4.11650e-01,  ..., -7.17188e+00, -7.16679e+00, -7.41265e+00],\n",
       "            ...,\n",
       "            [ 1.49597e-01, -3.89298e-01, -2.46371e-01,  ..., -6.58551e+00, -6.81194e+00, -7.59113e+00],\n",
       "            [ 1.25853e-01, -7.61917e-01,  2.55573e-01,  ..., -5.86286e+00, -5.77236e+00, -7.73055e+00],\n",
       "            [ 3.18480e-01,  1.43286e-01, -1.35235e-01,  ..., -3.78922e+00, -3.82270e+00, -6.50505e+00]],\n",
       " \n",
       "           [[-4.67868e-01,  9.68103e-02, -6.68329e-01,  ..., -6.01422e+00, -6.06895e+00, -7.02362e+00],\n",
       "            [ 2.45346e-02, -3.26566e-01, -8.28747e-01,  ..., -9.13249e+00, -7.72131e+00, -7.87114e+00],\n",
       "            [-1.31105e-01, -8.70810e-01, -6.83799e-01,  ..., -1.02013e+01, -8.10416e+00, -8.15806e+00],\n",
       "            ...,\n",
       "            [ 2.21245e-01, -3.44060e-01, -5.96235e-01,  ..., -9.46268e+00, -8.34107e+00, -8.26838e+00],\n",
       "            [-9.48046e-01,  3.23311e-01, -5.60571e-01,  ..., -9.33356e+00, -7.97358e+00, -8.48002e+00],\n",
       "            [ 2.77202e-01,  1.94120e-01, -5.32353e-01,  ..., -5.34957e+00, -5.15827e+00, -6.84345e+00]],\n",
       " \n",
       "           [[-4.10958e-01,  2.42855e-01, -6.81440e-01,  ..., -5.02189e+00, -5.87542e+00, -6.03359e+00],\n",
       "            [-2.92441e-01,  3.18814e-01, -6.60664e-01,  ..., -9.72233e+00, -7.88956e+00, -7.78910e+00],\n",
       "            [-8.08216e-01,  2.31449e-01, -6.20542e-01,  ..., -1.13929e+01, -9.26600e+00, -8.48023e+00],\n",
       "            ...,\n",
       "            [-1.08233e+00,  1.40251e-01, -5.22863e-01,  ..., -1.11120e+01, -9.46262e+00, -8.43857e+00],\n",
       "            [ 4.86934e-01,  3.41678e-01, -4.91409e-01,  ..., -8.83819e+00, -7.90552e+00, -7.89502e+00],\n",
       "            [-2.53421e-01,  4.77997e-03, -4.29960e-01,  ..., -6.20591e+00, -6.16837e+00, -6.87385e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-9.03000e-01,  2.41697e-01, -8.43498e-01,  ..., -7.55795e+00, -4.77543e+00, -7.60730e+00],\n",
       "            [ 8.49654e-01, -1.87725e-01, -6.25947e-01,  ..., -1.07795e+01, -8.39683e+00, -8.12195e+00],\n",
       "            [ 6.24663e-01, -8.57293e-01, -3.62576e-01,  ..., -1.10143e+01, -9.11621e+00, -7.82159e+00],\n",
       "            ...,\n",
       "            [-1.04331e+00,  1.10850e-01, -5.61554e-01,  ..., -1.15336e+01, -9.53375e+00, -8.51895e+00],\n",
       "            [-9.66584e-01, -5.57572e-01, -3.67658e-01,  ..., -1.06795e+01, -8.91183e+00, -8.47455e+00],\n",
       "            [ 7.37523e-01,  6.80919e-01, -4.42232e-01,  ..., -6.62037e+00, -5.58052e+00, -7.98964e+00]],\n",
       " \n",
       "           [[-2.67002e-02,  6.17616e-01, -7.29462e-01,  ..., -7.68495e+00, -6.18564e+00, -7.54069e+00],\n",
       "            [ 3.60228e-02, -3.18809e-02, -6.41215e-01,  ..., -9.51383e+00, -7.34739e+00, -7.51982e+00],\n",
       "            [ 2.78784e-01, -1.00208e+00, -5.85678e-01,  ..., -1.09305e+01, -7.98270e+00, -7.91393e+00],\n",
       "            ...,\n",
       "            [-4.99135e-01, -7.77025e-01,  6.32395e-03,  ..., -1.08448e+01, -8.45416e+00, -8.02954e+00],\n",
       "            [-1.18794e-02, -6.85096e-01,  1.72937e-01,  ..., -9.70193e+00, -7.68798e+00, -7.65591e+00],\n",
       "            [ 3.50175e-01,  6.48976e-01, -1.42952e-01,  ..., -6.23781e+00, -5.40870e+00, -7.18128e+00]],\n",
       " \n",
       "           [[-7.33834e-02,  3.67883e-01, -2.29226e-01,  ..., -4.42930e+00, -4.82669e+00, -5.57406e+00],\n",
       "            [ 1.47690e-02,  5.18661e-01, -2.35195e-01,  ..., -5.59237e+00, -5.67920e+00, -6.48677e+00],\n",
       "            [-3.88057e-01,  1.95528e-01, -4.68465e-01,  ..., -7.54841e+00, -6.67986e+00, -7.24097e+00],\n",
       "            ...,\n",
       "            [-8.12767e-01,  6.29450e-01, -6.57367e-02,  ..., -5.78285e+00, -5.34947e+00, -6.48298e+00],\n",
       "            [-5.67055e-01,  3.61528e-01, -3.27815e-02,  ..., -5.60615e+00, -5.65750e+00, -6.26799e+00],\n",
       "            [-5.63456e-02,  3.20577e-01,  8.11090e-02,  ..., -3.46963e+00, -4.57287e+00, -5.63440e+00]]],\n",
       " \n",
       " \n",
       "          [[[-8.60371e-02, -7.92923e-02, -5.85372e-01,  ..., -4.73824e+00, -3.09182e+00, -4.28531e+00],\n",
       "            [-1.12467e+00, -3.51111e-01, -2.24317e-01,  ..., -5.03009e+00, -2.76020e+00, -5.13560e+00],\n",
       "            [ 2.99525e-01,  8.33935e-02, -4.25821e-01,  ..., -5.70105e+00, -4.23027e+00, -5.71421e+00],\n",
       "            ...,\n",
       "            [ 1.70824e-01, -4.39289e-01, -3.58741e-01,  ..., -5.29601e+00, -3.99432e+00, -5.70017e+00],\n",
       "            [ 1.27990e-01, -7.91000e-01, -1.12771e-01,  ..., -4.98356e+00, -2.84335e+00, -5.72134e+00],\n",
       "            [ 2.42081e-01,  1.08009e-01, -4.75943e-01,  ..., -4.13734e+00, -2.21771e+00, -4.59058e+00]],\n",
       " \n",
       "           [[-3.35127e-01, -1.82518e-02, -7.96461e-01,  ..., -4.94455e+00, -3.79371e+00, -4.67221e+00],\n",
       "            [ 7.14469e-02, -4.46966e-01, -7.10243e-01,  ..., -6.59719e+00, -4.65007e+00, -6.29875e+00],\n",
       "            [-7.53697e-02, -8.91104e-01, -6.62942e-01,  ..., -6.75008e+00, -5.13280e+00, -6.52920e+00],\n",
       "            ...,\n",
       "            [ 2.50998e-01, -4.58230e-01, -5.21526e-01,  ..., -6.79897e+00, -4.16164e+00, -7.10570e+00],\n",
       "            [-9.38028e-01,  3.15594e-01, -6.19603e-01,  ..., -6.33238e+00, -4.80224e+00, -6.46101e+00],\n",
       "            [ 1.90264e-01,  1.60727e-01, -6.74735e-01,  ..., -4.83429e+00, -3.04963e+00, -4.79525e+00]],\n",
       " \n",
       "           [[-2.77173e-01,  1.21058e-01, -6.95702e-01,  ..., -5.11969e+00, -3.37714e+00, -4.00055e+00],\n",
       "            [-2.48339e-01,  2.88984e-01, -6.72909e-01,  ..., -7.22449e+00, -4.58168e+00, -6.26288e+00],\n",
       "            [-7.56938e-01,  2.47160e-01, -5.66172e-01,  ..., -8.10455e+00, -5.73535e+00, -6.89801e+00],\n",
       "            ...,\n",
       "            [-1.02688e+00,  3.46557e-02, -4.73552e-01,  ..., -7.95546e+00, -5.45082e+00, -7.00207e+00],\n",
       "            [ 4.79364e-01,  2.66618e-01, -3.97200e-01,  ..., -6.74121e+00, -4.55664e+00, -6.51718e+00],\n",
       "            [-3.34891e-01, -3.81756e-02, -5.56742e-01,  ..., -5.27786e+00, -3.65364e+00, -5.04372e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-7.48717e-01,  1.46352e-01, -9.59944e-01,  ..., -5.83053e+00, -2.96677e+00, -5.43640e+00],\n",
       "            [ 9.02113e-01, -2.95604e-01, -6.42298e-01,  ..., -7.31508e+00, -4.78419e+00, -6.39120e+00],\n",
       "            [ 6.43406e-01, -8.61298e-01, -4.84071e-01,  ..., -7.89323e+00, -5.57784e+00, -6.42421e+00],\n",
       "            ...,\n",
       "            [-9.85796e-01,  3.95502e-02, -6.32604e-01,  ..., -7.54403e+00, -5.43137e+00, -6.72892e+00],\n",
       "            [-9.50386e-01, -7.15261e-01, -3.34690e-01,  ..., -7.42454e+00, -4.57201e+00, -6.81979e+00],\n",
       "            [ 6.13301e-01,  6.33187e-01, -5.95568e-01,  ..., -5.25337e+00, -2.92365e+00, -5.60611e+00]],\n",
       " \n",
       "           [[ 1.19289e-01,  4.86794e-01, -7.56958e-01,  ..., -5.78522e+00, -3.20290e+00, -5.50376e+00],\n",
       "            [ 1.17694e-01, -2.07256e-01, -5.53289e-01,  ..., -6.63405e+00, -3.29675e+00, -6.12676e+00],\n",
       "            [ 3.07382e-01, -1.03843e+00, -5.47870e-01,  ..., -7.56152e+00, -4.81724e+00, -6.60747e+00],\n",
       "            ...,\n",
       "            [-4.33189e-01, -8.34786e-01, -1.04040e-01,  ..., -7.48849e+00, -4.09875e+00, -6.82705e+00],\n",
       "            [-1.22348e-02, -7.68148e-01, -1.46909e-02,  ..., -7.01686e+00, -3.82617e+00, -6.51272e+00],\n",
       "            [ 2.14937e-01,  6.03336e-01, -3.88539e-01,  ..., -5.38206e+00, -3.08042e+00, -5.55395e+00]],\n",
       " \n",
       "           [[ 3.57797e-02,  2.71361e-01, -5.55274e-01,  ..., -4.60773e+00, -3.66721e+00, -3.65215e+00],\n",
       "            [ 2.84679e-02,  4.13914e-01, -4.46833e-01,  ..., -4.94426e+00, -3.96014e+00, -4.73317e+00],\n",
       "            [-3.49589e-01,  9.58721e-02, -5.06214e-01,  ..., -5.87537e+00, -4.59131e+00, -5.72115e+00],\n",
       "            ...,\n",
       "            [-7.81144e-01,  5.33152e-01, -2.08942e-01,  ..., -5.09421e+00, -3.98459e+00, -4.63194e+00],\n",
       "            [-5.58426e-01,  2.69940e-01, -2.19376e-01,  ..., -5.09045e+00, -4.03432e+00, -4.60372e+00],\n",
       "            [-1.67765e-01,  2.38200e-01, -3.87894e-01,  ..., -4.15389e+00, -3.65103e+00, -3.34714e+00]]]]], device='cuda:0'))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_model(torch.rand((1,3,384,672)).cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8cea66b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calibrating model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|███▊                                                                                                                       | 1/32 [00:00<00:03,  7.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  6%|███████▋                                                                                                                   | 2/32 [00:00<00:03,  7.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 12%|███████████████▍                                                                                                           | 4/32 [00:00<00:02, 10.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 16%|███████████████████▏                                                                                                       | 5/32 [00:00<00:03,  7.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██████████████████████████████████▌                                                                                        | 9/32 [00:01<00:02,  9.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|█████████████████████████████████████████████▊                                                                            | 12/32 [00:01<00:02,  9.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 44%|█████████████████████████████████████████████████████▍                                                                    | 14/32 [00:01<00:01, 10.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 56%|████████████████████████████████████████████████████████████████████▋                                                     | 18/32 [00:01<00:01, 12.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 62%|████████████████████████████████████████████████████████████████████████████▎                                             | 20/32 [00:02<00:01,  9.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████████████████████████████████████████████████████████████████████████████████████████▌                              | 24/32 [00:02<00:00, 12.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 81%|███████████████████████████████████████████████████████████████████████████████████████████████████▏                      | 26/32 [00:02<00:00, 11.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 88%|██████████████████████████████████████████████████████████████████████████████████████████████████████████▊               | 28/32 [00:02<00:00, 11.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:03<00:00, 10.29it/s]\n",
      "W0615 14:21:07.297658 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.297943 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.298205 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.298437 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.298690 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.298940 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.299187 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.299429 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.299671 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0615 14:21:07.299917 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.300198 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.300467 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.300720 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.301927 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.302227 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.302485 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.302745 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.302993 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.303253 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.303500 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.303754 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.304003 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.304255 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.304504 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.304763 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.305017 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.305270 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.305521 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.305775 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.306021 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.306272 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.306515 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.306767 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.307014 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.307266 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.307508 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.307761 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.308007 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.308260 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.308503 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.308773 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.309014 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.309266 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.309506 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.309748 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.309993 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.310239 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.310478 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.310733 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.310975 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.311225 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.311465 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.311708 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.311947 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.312194 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.312433 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.312681 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.312966 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.313230 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.313497 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.313787 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.314025 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.314267 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.314509 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.314756 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.314993 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.315241 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.315475 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.315722 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.315958 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.316205 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.316444 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.316684 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.316970 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.317239 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.317495 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.317740 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.317979 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.318224 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.318461 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.318707 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.318972 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.319226 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.319468 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.319758 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.320005 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.320309 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.320715 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.320979 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.321226 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.321482 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.321730 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.321981 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.322227 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.322476 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.322720 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.322972 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.323214 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.323484 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.323761 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.324046 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.324327 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.324613 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.324901 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.325189 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.325469 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0615 14:21:07.325731 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.325976 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.326229 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.326488 140418879964992 tensor_quantizer.py:173] Disable MaxCalibrator\n",
      "W0615 14:21:07.326907 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.333152 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.333743 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.334409 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.334939 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.335577 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.336089 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.336731 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.337244 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.337879 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.338388 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.339020 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.339543 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.340261 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.340848 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.341699 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.342300 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.343064 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.343588 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.344222 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.344751 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.345417 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.345974 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.346717 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.347289 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.348095 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.348620 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.349313 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.349856 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.350526 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.351064 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.351716 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.352240 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.352968 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.353541 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.354359 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.354936 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.355613 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.356154 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.356828 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.357367 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.358026 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.358580 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.359307 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([512, 1, 1, 1]).\n",
      "W0615 14:21:07.359817 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.360447 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.360973 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.361639 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.362296 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.363449 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.364028 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.364446 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.365420 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.366088 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.366628 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.367357 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.367899 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.368629 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.369445 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.369867 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.370381 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.371045 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.371552 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.372266 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.372836 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.373553 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.374063 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.374701 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.375265 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.376125 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0615 14:21:07.376688 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.377484 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.378035 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.378716 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.379264 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.380019 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.380566 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.381316 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([32, 1, 1, 1]).\n",
      "W0615 14:21:07.381878 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.382623 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.383166 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.383842 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.384402 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.385148 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.385667 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.386374 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.386914 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.387588 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.388133 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.388853 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([64, 1, 1, 1]).\n",
      "W0615 14:21:07.389415 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.390164 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.390745 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.391426 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.392004 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.392730 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.393247 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.393894 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.394410 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.395063 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.395580 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.396226 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.396762 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.397490 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.398006 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.398657 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([128, 1, 1, 1]).\n",
      "W0615 14:21:07.399177 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.399824 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([256, 1, 1, 1]).\n",
      "W0615 14:21:07.400348 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([]).\n",
      "W0615 14:21:07.401007 140418879964992 tensor_quantizer.py:237] Load calibrated amax, shape=torch.Size([512, 1, 1, 1]).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skip quantization: \n",
      "Skip quantization: model\n",
      "Skip quantization: model.0\n",
      "Skip quantization: model.0.conv\n",
      "model.0.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=1.0000 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.0.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.8467, 9.8950](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.0.act\n",
      "Skip quantization: model.1\n",
      "Skip quantization: model.1.conv\n",
      "model.1.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=21.9565 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.1.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.1028, 1.0065](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.1.act\n",
      "Skip quantization: model.2\n",
      "Skip quantization: model.2.conv\n",
      "model.2.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=31.6012 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.2.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.3582, 1.2886](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.2.act\n",
      "Skip quantization: model.3\n",
      "Skip quantization: model.3.conv\n",
      "model.3.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=31.6012 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.3.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.3486, 1.3525](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.3.act\n",
      "Skip quantization: model.4\n",
      "Skip quantization: model.4.conv\n",
      "model.4.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=10.4433 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.4.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.1656, 0.7424](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.4.act\n",
      "Skip quantization: model.5\n",
      "Skip quantization: model.5.conv\n",
      "model.5.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=17.1036 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.5.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.1998, 0.4605](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.5.act\n",
      "Skip quantization: model.6\n",
      "Skip quantization: model.7\n",
      "Skip quantization: model.7.conv\n",
      "model.7.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=17.1036 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.7.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.2378, 0.6717](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.7.act\n",
      "Skip quantization: model.8\n",
      "Skip quantization: model.8.m\n",
      "Skip quantization: model.9\n",
      "Skip quantization: model.9.conv\n",
      "model.9.conv._input_quantizer           : TensorQuantizer(8bit fake per-tensor amax=11.7819 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.9.conv._weight_quantizer          : TensorQuantizer(8bit fake axis=0 amax=[0.1816, 0.7992](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.9.act\n",
      "Skip quantization: model.10\n",
      "Skip quantization: model.10.conv\n",
      "model.10.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.7819 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.10.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1792, 0.7680](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.10.act\n",
      "Skip quantization: model.11\n",
      "Skip quantization: model.11.conv\n",
      "model.11.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.7092 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.11.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1380, 0.3033](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.11.act\n",
      "Skip quantization: model.12\n",
      "Skip quantization: model.12.conv\n",
      "model.12.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.3041 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.12.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1401, 0.3056](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.12.act\n",
      "Skip quantization: model.13\n",
      "Skip quantization: model.14\n",
      "Skip quantization: model.14.conv\n",
      "model.14.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=13.7202 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.14.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1422, 0.5593](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.14.act\n",
      "Skip quantization: model.15\n",
      "Skip quantization: model.15.m\n",
      "Skip quantization: model.16\n",
      "Skip quantization: model.16.conv\n",
      "model.16.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=12.3185 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.16.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1325, 0.5530](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.16.act\n",
      "Skip quantization: model.17\n",
      "Skip quantization: model.17.conv\n",
      "model.17.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=12.3185 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.17.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1662, 0.5371](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.17.act\n",
      "Skip quantization: model.18\n",
      "Skip quantization: model.18.conv\n",
      "model.18.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.5341 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.18.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0773, 0.2254](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.18.act\n",
      "Skip quantization: model.19\n",
      "Skip quantization: model.19.conv\n",
      "model.19.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.2964 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.19.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0871, 0.2030](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.19.act\n",
      "Skip quantization: model.20\n",
      "Skip quantization: model.21\n",
      "Skip quantization: model.21.conv\n",
      "model.21.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.4538 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.21.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1095, 0.7497](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.21.act\n",
      "Skip quantization: model.22\n",
      "Skip quantization: model.22.m\n",
      "Skip quantization: model.23\n",
      "Skip quantization: model.23.conv\n",
      "model.23.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.7248 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.23.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0860, 0.2745](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.23.act\n",
      "Skip quantization: model.24\n",
      "Skip quantization: model.24.conv\n",
      "model.24.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.7248 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.24.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1073, 0.5097](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.24.act\n",
      "Skip quantization: model.25\n",
      "Skip quantization: model.25.conv\n",
      "model.25.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.7734 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.25.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0468, 0.1695](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.25.act\n",
      "Skip quantization: model.26\n",
      "Skip quantization: model.26.conv\n",
      "model.26.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.2573 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.26.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0491, 0.1701](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.26.act\n",
      "Skip quantization: model.27\n",
      "Skip quantization: model.28\n",
      "Skip quantization: model.28.conv\n",
      "model.28.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.3556 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.28.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0789, 0.1798](512) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.28.act\n",
      "Skip quantization: model.29\n",
      "Skip quantization: model.29.conv\n",
      "model.29.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=6.8431 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.29.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1131, 0.2692](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.29.act\n",
      "Skip quantization: model.30\n",
      "Skip quantization: model.30.conv\n",
      "model.30.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=6.8431 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.30.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0836, 0.2447](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.30.act\n",
      "Skip quantization: model.31\n",
      "Skip quantization: model.31.m\n",
      "Skip quantization: model.32\n",
      "Skip quantization: model.32.m\n",
      "Skip quantization: model.33\n",
      "Skip quantization: model.33.m\n",
      "Skip quantization: model.34\n",
      "Skip quantization: model.35\n",
      "Skip quantization: model.35.conv\n",
      "model.35.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=6.9201 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.35.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0546, 0.1576](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.35.act\n",
      "Skip quantization: model.36\n",
      "Skip quantization: model.37\n",
      "Skip quantization: model.37.conv\n",
      "model.37.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=5.8761 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.37.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1058, 0.3790](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.37.act\n",
      "Skip quantization: model.38\n",
      "Skip quantization: model.38.conv\n",
      "model.38.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=5.9318 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.38.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1772, 0.4712](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.38.act\n",
      "Skip quantization: model.39\n",
      "Skip quantization: model.40\n",
      "Skip quantization: model.40.conv\n",
      "model.40.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.7248 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.40.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1236, 0.3978](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.40.act\n",
      "Skip quantization: model.41\n",
      "Skip quantization: model.42\n",
      "Skip quantization: model.42.conv\n",
      "model.42.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.6507 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.42.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1526, 0.4597](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.42.act\n",
      "Skip quantization: model.43\n",
      "Skip quantization: model.43.conv\n",
      "model.43.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.6507 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.43.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1507, 0.3301](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.43.act\n",
      "Skip quantization: model.44\n",
      "Skip quantization: model.44.conv\n",
      "model.44.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=7.6914 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.44.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0871, 0.3623](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.44.act\n",
      "Skip quantization: model.45\n",
      "Skip quantization: model.45.conv\n",
      "model.45.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=7.5127 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.45.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0984, 0.2614](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.45.act\n",
      "Skip quantization: model.46\n",
      "Skip quantization: model.47\n",
      "Skip quantization: model.47.conv\n",
      "model.47.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.4010 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.47.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1394, 0.5261](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.47.act\n",
      "Skip quantization: model.48\n",
      "Skip quantization: model.48.conv\n",
      "model.48.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.0094 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.48.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.2262, 0.8891](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.48.act\n",
      "Skip quantization: model.49\n",
      "Skip quantization: model.50\n",
      "Skip quantization: model.50.conv\n",
      "model.50.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=12.3185 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.50.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1668, 0.8992](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.50.act\n",
      "Skip quantization: model.51\n",
      "Skip quantization: model.52\n",
      "Skip quantization: model.52.conv\n",
      "model.52.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.8600 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.52.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.2115, 0.6500](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.52.act\n",
      "Skip quantization: model.53\n",
      "Skip quantization: model.53.conv\n",
      "model.53.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.8600 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.53.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1715, 0.4547](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.53.act\n",
      "Skip quantization: model.54\n",
      "Skip quantization: model.54.conv\n",
      "model.54.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=7.2781 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.54.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.2098, 0.6950](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.54.act\n",
      "Skip quantization: model.55\n",
      "Skip quantization: model.55.conv\n",
      "model.55.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.7555 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.55.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.2554, 0.8387](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.55.act\n",
      "Skip quantization: model.56\n",
      "Skip quantization: model.57\n",
      "Skip quantization: model.57.conv\n",
      "model.57.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.4869 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.57.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1830, 0.7702](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.57.act\n",
      "Skip quantization: model.58\n",
      "Skip quantization: model.58.conv\n",
      "model.58.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.5577 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.58.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1000, 0.3975](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.58.act\n",
      "Skip quantization: model.59\n",
      "Skip quantization: model.60\n",
      "Skip quantization: model.60.conv\n",
      "model.60.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.9869 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.60.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1485, 0.5582](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.60.act\n",
      "Skip quantization: model.61\n",
      "Skip quantization: model.61.conv\n",
      "model.61.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.9869 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.61.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1336, 0.6178](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.61.act\n",
      "Skip quantization: model.62\n",
      "Skip quantization: model.62.conv\n",
      "model.62.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.7306 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.62.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1078, 0.5156](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.62.act\n",
      "Skip quantization: model.63\n",
      "Skip quantization: model.63.conv\n",
      "model.63.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=9.3281 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.63.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1375, 0.6402](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.63.act\n",
      "Skip quantization: model.64\n",
      "Skip quantization: model.65\n",
      "Skip quantization: model.65.conv\n",
      "model.65.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.4749 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.65.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1214, 0.9532](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.65.act\n",
      "Skip quantization: model.66\n",
      "Skip quantization: model.66.conv\n",
      "model.66.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.9742 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.66.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0747, 0.3336](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.66.act\n",
      "Skip quantization: model.67\n",
      "Skip quantization: model.68\n",
      "Skip quantization: model.68.conv\n",
      "model.68.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.2716 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.68.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0944, 0.3416](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.68.act\n",
      "Skip quantization: model.69\n",
      "Skip quantization: model.69.conv\n",
      "model.69.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=10.2716 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.69.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0846, 0.4276](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.69.act\n",
      "Skip quantization: model.70\n",
      "Skip quantization: model.70.conv\n",
      "model.70.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=12.6105 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.70.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0681, 0.4963](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.70.act\n",
      "Skip quantization: model.71\n",
      "Skip quantization: model.71.conv\n",
      "model.71.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.0984 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.71.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0713, 0.5777](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.71.act\n",
      "Skip quantization: model.72\n",
      "Skip quantization: model.73\n",
      "Skip quantization: model.73.conv\n",
      "model.73.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=12.6105 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.73.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0650, 0.7201](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.73.act\n",
      "Skip quantization: model.74\n",
      "Skip quantization: model.74.conv\n",
      "model.74.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.5577 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.74.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.1087, 1.2581](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.74.act\n",
      "Skip quantization: model.75\n",
      "Skip quantization: model.75.conv\n",
      "model.75.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=11.9742 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.75.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0504, 1.1487](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.75.act\n",
      "Skip quantization: model.76\n",
      "Skip quantization: model.76.conv\n",
      "model.76.conv._input_quantizer          : TensorQuantizer(8bit fake per-tensor amax=8.8608 calibrator=MaxCalibrator scale=1.0 quant)\n",
      "model.76.conv._weight_quantizer         : TensorQuantizer(8bit fake axis=0 amax=[0.0312, 0.6864](512) calibrator=MaxCalibrator scale=1.0 quant)\n",
      "Skip quantization: model.76.act\n",
      "Skip quantization: model.77\n",
      "Skip quantization: model.77.m\n",
      "Skip quantization: model.77.m.0\n",
      "Skip quantization: model.77.m.1\n",
      "Skip quantization: model.77.m.2\n",
      "Skip quantization: model.77.ia\n",
      "Skip quantization: model.77.ia.0\n",
      "Skip quantization: model.77.ia.1\n",
      "Skip quantization: model.77.ia.2\n",
      "Skip quantization: model.77.im\n",
      "Skip quantization: model.77.im.0\n",
      "Skip quantization: model.77.im.1\n",
      "Skip quantization: model.77.im.2\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    calibrate_model(\n",
    "        model=qat_model,\n",
    "        model_name='yolov7-tiny',\n",
    "        data_loader=dataloader,\n",
    "        num_calib_batch=32,\n",
    "        calibrator=\"max\",\n",
    "        hist_percentile=[99.9, 99.99, 99.999, 99.9999],\n",
    "        out_dir=\"./runs/train/yolov7-tiny/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8aa49cd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[[[[-4.80892e-01,  2.16270e-01,  4.86674e-01,  ..., -5.82189e+00, -5.50311e+00, -2.94533e+00],\n",
       "            [-6.06903e-01, -4.96320e-02,  5.14759e-01,  ..., -7.16952e+00, -5.98059e+00, -4.31843e+00],\n",
       "            [-7.81228e-01,  9.20628e-02,  4.61599e-01,  ..., -7.55950e+00, -6.04208e+00, -4.39546e+00],\n",
       "            ...,\n",
       "            [-8.87838e-02, -1.49041e-01,  8.43424e-01,  ..., -7.36928e+00, -6.51112e+00, -5.47578e+00],\n",
       "            [-1.29246e+00, -1.56100e-01,  1.27963e+00,  ..., -6.00930e+00, -5.14863e+00, -3.37378e+00],\n",
       "            [-3.33980e-01,  1.59899e-01,  1.14406e+00,  ..., -4.98250e+00, -4.58480e+00, -2.78846e+00]],\n",
       " \n",
       "           [[-3.28819e-01, -5.15180e-02,  5.64076e-01,  ..., -5.64666e+00, -5.58934e+00, -3.09178e+00],\n",
       "            [-9.41898e-01, -8.29980e-01,  6.91706e-01,  ..., -5.99931e+00, -4.94950e+00, -3.50759e+00],\n",
       "            [-2.51756e-01, -1.09771e+00,  5.05238e-01,  ..., -7.64555e+00, -5.64651e+00, -4.33606e+00],\n",
       "            ...,\n",
       "            [-3.04189e-01, -1.01791e+00,  9.07576e-01,  ..., -6.80487e+00, -5.44767e+00, -4.22387e+00],\n",
       "            [-9.46548e-02, -1.37691e-01,  1.11417e+00,  ..., -6.36372e+00, -5.15663e+00, -3.40346e+00],\n",
       "            [ 2.24452e-01, -8.43449e-01,  6.83576e-01,  ..., -5.45801e+00, -5.68650e+00, -3.14355e+00]],\n",
       " \n",
       "           [[-3.25946e-01, -3.39069e-01,  3.48745e-01,  ..., -7.43750e+00, -6.30436e+00, -5.69675e+00],\n",
       "            [-9.90131e-01,  3.27159e-01,  4.34550e-01,  ..., -7.23527e+00, -5.68685e+00, -4.64723e+00],\n",
       "            [-4.74711e-01,  4.66489e-02,  4.76744e-01,  ..., -7.78721e+00, -5.58552e+00, -4.83307e+00],\n",
       "            ...,\n",
       "            [-3.17674e-01, -1.40070e-01,  5.98650e-01,  ..., -7.16251e+00, -5.21726e+00, -4.74738e+00],\n",
       "            [-3.11820e-01,  2.79884e-01,  7.14975e-01,  ..., -7.21343e+00, -5.63209e+00, -5.09741e+00],\n",
       "            [-1.42362e-01, -7.08733e-01,  1.01042e+00,  ..., -5.41688e+00, -4.87847e+00, -3.03818e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-2.35683e-01, -7.15734e-01,  3.75932e-01,  ..., -5.94345e+00, -5.30138e+00, -3.56979e+00],\n",
       "            [-5.69279e-01,  4.06133e-01,  8.77030e-01,  ..., -6.47570e+00, -4.77502e+00, -4.85989e+00],\n",
       "            [-7.04618e-01, -7.20657e-02,  1.11650e+00,  ..., -6.85725e+00, -4.50108e+00, -4.40791e+00],\n",
       "            ...,\n",
       "            [-6.00052e-01, -2.41996e-01,  7.56743e-01,  ..., -7.33519e+00, -5.50936e+00, -4.55808e+00],\n",
       "            [-7.89088e-01,  8.82362e-03,  9.97613e-01,  ..., -6.75395e+00, -5.12750e+00, -3.70682e+00],\n",
       "            [-3.08638e-01,  5.19888e-01,  5.46282e-01,  ..., -5.87470e+00, -4.90076e+00, -2.96358e+00]],\n",
       " \n",
       "           [[-1.72499e-01, -1.51399e-01,  5.49669e-01,  ..., -6.62426e+00, -5.86237e+00, -4.16346e+00],\n",
       "            [-6.60497e-01,  2.48144e-01,  8.30787e-01,  ..., -6.94169e+00, -5.63343e+00, -5.20754e+00],\n",
       "            [-1.08220e+00,  6.07740e-01,  1.15646e+00,  ..., -6.30040e+00, -4.57185e+00, -4.10208e+00],\n",
       "            ...,\n",
       "            [-6.12322e-01,  8.60948e-01,  8.10537e-01,  ..., -6.77423e+00, -5.53438e+00, -4.38401e+00],\n",
       "            [-4.18058e-01,  8.89066e-01,  1.04111e+00,  ..., -6.95435e+00, -5.36617e+00, -4.01161e+00],\n",
       "            [-1.34934e-01,  2.02298e-01,  6.90016e-01,  ..., -6.85771e+00, -6.13576e+00, -4.03243e+00]],\n",
       " \n",
       "           [[-4.17911e-01, -8.65882e-01,  7.03296e-01,  ..., -5.46189e+00, -5.29011e+00, -2.11867e+00],\n",
       "            [-4.24667e-01, -1.38112e-01,  8.19158e-01,  ..., -6.24139e+00, -5.49866e+00, -3.17930e+00],\n",
       "            [-4.18747e-01,  2.51755e-01,  8.06846e-01,  ..., -6.68910e+00, -5.61763e+00, -3.14041e+00],\n",
       "            ...,\n",
       "            [-4.86441e-01, -1.89531e-01,  8.97513e-01,  ..., -6.99466e+00, -6.20450e+00, -3.67860e+00],\n",
       "            [-7.09292e-01, -1.28774e-01,  8.73542e-01,  ..., -7.59818e+00, -6.59679e+00, -4.34193e+00],\n",
       "            [-5.46117e-01, -5.15540e-01,  1.27024e+00,  ..., -5.90942e+00, -5.71964e+00, -2.86071e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 2.31342e-01,  4.11293e-01,  4.30565e-01,  ..., -5.69661e+00, -5.71731e+00, -1.62738e+00],\n",
       "            [ 5.68930e-01,  8.91260e-02,  2.70610e-01,  ..., -7.92313e+00, -6.55770e+00, -3.15526e+00],\n",
       "            [ 6.22616e-01,  2.10665e-01,  1.93821e-01,  ..., -8.59395e+00, -6.40848e+00, -3.66904e+00],\n",
       "            ...,\n",
       "            [ 2.80114e-01, -6.84372e-03,  5.33147e-01,  ..., -8.14287e+00, -6.85230e+00, -4.50818e+00],\n",
       "            [ 1.17073e-01,  7.43696e-02,  1.24156e+00,  ..., -5.97985e+00, -5.44079e+00, -2.56164e+00],\n",
       "            [-5.05371e-02,  5.80151e-01,  8.10292e-01,  ..., -4.41617e+00, -4.43977e+00, -2.16928e+00]],\n",
       " \n",
       "           [[-4.10572e-02,  4.54807e-01,  5.49102e-01,  ..., -6.40926e+00, -5.92197e+00, -1.89308e+00],\n",
       "            [ 4.28303e-01, -7.85692e-01,  7.43645e-01,  ..., -6.19454e+00, -5.40445e+00, -3.04727e+00],\n",
       "            [ 1.31731e+00, -1.08700e+00,  4.99029e-01,  ..., -8.58007e+00, -5.98161e+00, -4.39941e+00],\n",
       "            ...,\n",
       "            [ 4.95884e-01, -7.38782e-01,  8.97513e-01,  ..., -7.40168e+00, -5.68614e+00, -3.72335e+00],\n",
       "            [ 1.31806e+00, -5.51388e-02,  1.41440e+00,  ..., -6.66053e+00, -5.44098e+00, -2.60265e+00],\n",
       "            [ 1.08429e-01, -5.01564e-01,  3.86070e-01,  ..., -5.51179e+00, -5.68856e+00, -2.31142e+00]],\n",
       " \n",
       "           [[-3.24836e-01, -1.39561e-01,  2.84023e-01,  ..., -8.32240e+00, -6.36806e+00, -4.57645e+00],\n",
       "            [ 5.22960e-01,  3.46178e-01,  3.17755e-01,  ..., -8.41854e+00, -6.07564e+00, -4.75925e+00],\n",
       "            [ 9.95789e-01,  9.72370e-02,  3.53232e-01,  ..., -8.92392e+00, -5.85077e+00, -4.83119e+00],\n",
       "            ...,\n",
       "            [ 6.98842e-01,  7.23070e-02,  4.94596e-01,  ..., -8.03669e+00, -5.36313e+00, -4.49228e+00],\n",
       "            [ 7.19825e-01,  3.96264e-01,  7.15577e-01,  ..., -8.22183e+00, -5.85080e+00, -4.35363e+00],\n",
       "            [ 1.84216e-01, -3.93739e-01,  6.97157e-01,  ..., -5.02001e+00, -4.64357e+00, -2.58338e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.69365e-01, -6.13121e-02,  2.88462e-01,  ..., -6.40740e+00, -5.27621e+00, -3.24580e+00],\n",
       "            [ 3.56085e-01,  4.81533e-01,  9.77316e-01,  ..., -7.04088e+00, -5.07656e+00, -4.58676e+00],\n",
       "            [ 6.54926e-01, -4.12570e-02,  1.29601e+00,  ..., -7.43510e+00, -4.59012e+00, -4.07000e+00],\n",
       "            ...,\n",
       "            [ 2.55746e-01, -8.58382e-02,  7.67224e-01,  ..., -8.35414e+00, -5.69374e+00, -4.51217e+00],\n",
       "            [ 4.53124e-01,  3.28400e-01,  1.13599e+00,  ..., -7.47062e+00, -5.13417e+00, -3.53733e+00],\n",
       "            [ 2.76877e-01,  8.61841e-01,  3.60970e-01,  ..., -6.10047e+00, -4.87509e+00, -2.55201e+00]],\n",
       " \n",
       "           [[-1.30849e-01,  1.62477e-01,  4.25952e-01,  ..., -7.11675e+00, -5.81909e+00, -3.89556e+00],\n",
       "            [ 1.00341e-01,  6.36387e-01,  7.69757e-01,  ..., -7.72244e+00, -5.54348e+00, -4.91185e+00],\n",
       "            [ 5.32341e-01,  9.61115e-01,  1.21241e+00,  ..., -7.06244e+00, -4.54815e+00, -3.83563e+00],\n",
       "            ...,\n",
       "            [ 3.38026e-01,  1.41889e+00,  7.90727e-01,  ..., -7.67851e+00, -5.55085e+00, -3.91631e+00],\n",
       "            [ 1.27259e+00,  1.14188e+00,  1.05952e+00,  ..., -7.86894e+00, -5.42215e+00, -3.68789e+00],\n",
       "            [-6.61479e-02,  7.26086e-01,  4.43065e-01,  ..., -7.43132e+00, -6.06099e+00, -3.57719e+00]],\n",
       " \n",
       "           [[ 1.74061e-01, -7.06992e-01,  5.93693e-01,  ..., -5.45000e+00, -5.13360e+00, -2.53099e+00],\n",
       "            [ 5.57224e-01, -1.61753e-01,  6.19512e-01,  ..., -6.77903e+00, -5.50624e+00, -3.00839e+00],\n",
       "            [ 6.88685e-01,  4.55185e-03,  6.77970e-01,  ..., -7.42892e+00, -5.40447e+00, -2.86614e+00],\n",
       "            ...,\n",
       "            [ 5.47938e-01, -2.34730e-01,  6.06148e-01,  ..., -7.68995e+00, -6.41526e+00, -3.53516e+00],\n",
       "            [ 1.70632e-01, -1.88094e-01,  5.66089e-01,  ..., -8.50677e+00, -6.85867e+00, -4.23490e+00],\n",
       "            [-9.98779e-03, -4.36963e-01,  8.37824e-01,  ..., -5.73609e+00, -5.87368e+00, -2.88982e+00]]],\n",
       " \n",
       " \n",
       "          [[[-1.91045e-01,  6.01167e-01, -5.09683e-02,  ..., -4.22721e+00, -5.17744e+00, -7.76403e-01],\n",
       "            [-2.35563e-01,  1.01103e-01,  6.23386e-02,  ..., -6.32577e+00, -5.50530e+00, -2.05262e+00],\n",
       "            [-2.77643e-01,  1.42280e-01, -1.36790e-01,  ..., -7.16714e+00, -5.30765e+00, -3.06886e+00],\n",
       "            ...,\n",
       "            [ 8.41269e-02, -1.04338e-02,  3.53416e-02,  ..., -6.76476e+00, -6.18514e+00, -3.18370e+00],\n",
       "            [-4.93224e-01,  1.59555e-01,  7.55685e-01,  ..., -4.42818e+00, -4.60226e+00, -1.66079e+00],\n",
       "            [-5.00620e-01,  8.50793e-01,  3.31702e-01,  ..., -3.53203e+00, -4.28162e+00, -2.35645e+00]],\n",
       " \n",
       "           [[-1.10943e-01,  8.87727e-01, -6.68953e-02,  ..., -5.73965e+00, -5.37060e+00, -1.72486e+00],\n",
       "            [-4.99476e-01, -4.52270e-01,  4.10553e-01,  ..., -5.26371e+00, -4.49420e+00, -2.81314e+00],\n",
       "            [ 1.58477e-01, -7.07129e-01,  5.75275e-02,  ..., -7.99753e+00, -4.88841e+00, -4.51612e+00],\n",
       "            ...,\n",
       "            [-7.89974e-02, -3.57260e-01,  4.10507e-01,  ..., -6.51958e+00, -4.77922e+00, -3.11209e+00],\n",
       "            [ 5.22546e-01,  2.53806e-01,  8.64065e-01,  ..., -5.79366e+00, -4.52858e+00, -2.05439e+00],\n",
       "            [-9.95274e-02, -1.03205e-01,  2.68494e-02,  ..., -4.93707e+00, -5.40208e+00, -1.24371e+00]],\n",
       " \n",
       "           [[-2.69270e-01,  1.52171e-01, -2.93135e-01,  ..., -7.62982e+00, -5.90678e+00, -4.53092e+00],\n",
       "            [-3.72759e-01,  5.49802e-01, -7.79665e-02,  ..., -7.49037e+00, -5.01499e+00, -4.88489e+00],\n",
       "            [-5.93689e-02,  4.93213e-01, -1.50304e-01,  ..., -8.09095e+00, -4.78551e+00, -5.06035e+00],\n",
       "            ...,\n",
       "            [-4.67264e-02,  3.27727e-01,  4.24265e-02,  ..., -7.09546e+00, -4.45485e+00, -4.24519e+00],\n",
       "            [ 7.57091e-02,  5.75141e-01,  2.51450e-01,  ..., -7.36954e+00, -4.93639e+00, -3.87372e+00],\n",
       "            [-3.20643e-01, -1.20518e-01,  3.15797e-01,  ..., -4.19149e+00, -4.30538e+00, -2.10053e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.20621e-01,  6.30293e-01, -2.97156e-01,  ..., -5.97510e+00, -4.80450e+00, -3.19646e+00],\n",
       "            [-1.58135e-01,  6.49360e-01,  4.49424e-01,  ..., -6.26695e+00, -4.08564e+00, -4.54702e+00],\n",
       "            [-1.91767e-01,  2.64786e-01,  6.79546e-01,  ..., -6.83094e+00, -3.80726e+00, -4.74716e+00],\n",
       "            ...,\n",
       "            [-1.55865e-01,  2.50984e-01,  2.23764e-01,  ..., -7.98168e+00, -4.93211e+00, -5.55073e+00],\n",
       "            [-2.27238e-01,  6.28342e-01,  5.20406e-01,  ..., -6.86790e+00, -4.37686e+00, -4.24310e+00],\n",
       "            [-1.36375e-01,  1.13141e+00, -1.01739e-01,  ..., -5.37947e+00, -4.36015e+00, -2.36805e+00]],\n",
       " \n",
       "           [[-1.20721e-01,  5.72586e-01, -2.25421e-01,  ..., -6.23060e+00, -5.43713e+00, -3.55986e+00],\n",
       "            [-2.97759e-01,  9.96576e-01,  2.59492e-01,  ..., -6.87506e+00, -4.80344e+00, -5.28549e+00],\n",
       "            [-1.59202e-01,  1.44889e+00,  5.32891e-01,  ..., -6.55042e+00, -3.69974e+00, -4.56202e+00],\n",
       "            ...,\n",
       "            [-2.11332e-01,  1.97248e+00,  2.49453e-01,  ..., -6.92255e+00, -4.84864e+00, -4.17141e+00],\n",
       "            [ 2.54843e-01,  1.60323e+00,  3.24189e-01,  ..., -7.27934e+00, -4.70058e+00, -4.33929e+00],\n",
       "            [-1.75586e-01,  1.18136e+00, -1.72400e-01,  ..., -6.80504e+00, -5.72576e+00, -3.64271e+00]],\n",
       " \n",
       "           [[ 6.21668e-02, -3.11279e-01, -3.34550e-02,  ..., -4.93043e+00, -5.06972e+00, -2.34386e+00],\n",
       "            [-1.21233e-01, -5.26438e-02,  1.83721e-01,  ..., -6.08709e+00, -4.99897e+00, -2.74864e+00],\n",
       "            [ 6.21772e-02,  9.74414e-02,  1.64942e-01,  ..., -6.92769e+00, -5.01026e+00, -3.10241e+00],\n",
       "            ...,\n",
       "            [-3.03425e-02, -1.55967e-01, -2.41935e-02,  ..., -7.02531e+00, -5.94828e+00, -3.67427e+00],\n",
       "            [-3.51871e-01, -1.98497e-01, -9.49375e-02,  ..., -7.80895e+00, -6.44818e+00, -4.25268e+00],\n",
       "            [-4.59247e-01, -2.35425e-01,  1.79792e-01,  ..., -5.26103e+00, -5.85111e+00, -3.18215e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-3.66343e-01, -1.76690e-01,  1.32897e-02,  ..., -6.85739e+00, -4.25988e+00, -6.24257e+00],\n",
       "            [ 4.66276e-01, -2.03010e-01,  1.11599e-01,  ..., -9.58695e+00, -5.76774e+00, -6.94015e+00],\n",
       "            [-6.30949e-01,  3.65676e-03,  2.16214e-01,  ..., -9.74585e+00, -5.79796e+00, -6.85754e+00],\n",
       "            ...,\n",
       "            [-2.81506e-01, -5.30117e-01,  8.19732e-01,  ..., -6.05634e+00, -3.49203e+00, -4.96844e+00],\n",
       "            [ 4.98835e-01, -2.85640e-01,  1.18076e+00,  ..., -6.00457e+00, -3.61382e+00, -5.43938e+00],\n",
       "            [-1.30720e-01,  2.91228e-03,  1.07890e+00,  ..., -5.38645e+00, -2.40582e+00, -5.79188e+00]],\n",
       " \n",
       "           [[-4.83629e-01, -1.99903e-01, -9.13226e-02,  ..., -9.43628e+00, -4.97489e+00, -6.76669e+00],\n",
       "            [ 3.35618e-01,  1.23584e-01,  2.28139e-01,  ..., -1.16963e+01, -6.33535e+00, -7.06777e+00],\n",
       "            [-2.76549e-01,  2.76589e-01,  4.67789e-01,  ..., -1.06708e+01, -5.10057e+00, -6.44130e+00],\n",
       "            ...,\n",
       "            [ 1.53058e-01,  5.81170e-01,  4.47800e-01,  ..., -8.47259e+00, -5.29086e+00, -5.43439e+00],\n",
       "            [ 3.81079e-01, -4.84241e-01,  8.00847e-01,  ..., -7.37255e+00, -4.39228e+00, -4.79089e+00],\n",
       "            [ 7.46118e-02, -6.42692e-01,  6.37710e-01,  ..., -5.82071e+00, -3.48250e+00, -5.48283e+00]],\n",
       " \n",
       "           [[-3.14848e-01,  8.68502e-02,  1.61939e-01,  ..., -8.90615e+00, -4.67930e+00, -6.46801e+00],\n",
       "            [-4.76576e-01, -4.04632e-02,  2.51372e-01,  ..., -1.22003e+01, -6.28545e+00, -7.42603e+00],\n",
       "            [-6.97840e-01,  5.24016e-01,  4.39591e-01,  ..., -1.06499e+01, -5.40849e+00, -6.80276e+00],\n",
       "            ...,\n",
       "            [-5.06791e-01,  3.59850e-01,  5.08380e-01,  ..., -9.25983e+00, -5.22070e+00, -5.87820e+00],\n",
       "            [ 4.65401e-01,  7.45466e-01,  3.72992e-01,  ..., -9.02674e+00, -6.08081e+00, -6.16034e+00],\n",
       "            [ 1.66108e-01,  5.61313e-01,  4.25113e-01,  ..., -6.69862e+00, -3.18964e+00, -5.45648e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.66424e-01,  6.62825e-01, -7.05966e-03,  ..., -8.31714e+00, -3.95586e+00, -5.90568e+00],\n",
       "            [ 3.99745e-01, -6.92430e-01,  3.46471e-01,  ..., -1.07933e+01, -5.57736e+00, -6.57621e+00],\n",
       "            [ 2.54960e-01, -5.31596e-01,  4.39893e-01,  ..., -1.10682e+01, -5.40781e+00, -6.69323e+00],\n",
       "            ...,\n",
       "            [ 5.67764e-01,  4.46827e-01,  3.27459e-01,  ..., -9.19988e+00, -4.65336e+00, -5.47004e+00],\n",
       "            [ 1.16919e-01, -3.18450e-01,  6.93708e-02,  ..., -1.04961e+01, -5.31631e+00, -6.54850e+00],\n",
       "            [ 5.90247e-01, -3.26479e-01,  3.86301e-02,  ..., -7.44596e+00, -2.64124e+00, -5.57321e+00]],\n",
       " \n",
       "           [[-4.93325e-01,  5.10174e-02, -2.21195e-02,  ..., -8.48102e+00, -5.11506e+00, -6.37174e+00],\n",
       "            [-1.88103e-02,  1.06694e+00,  6.27909e-01,  ..., -9.60351e+00, -5.16246e+00, -6.20747e+00],\n",
       "            [ 4.99402e-01, -1.11211e-01,  5.95889e-01,  ..., -9.57621e+00, -4.79909e+00, -5.91427e+00],\n",
       "            ...,\n",
       "            [ 2.26260e-01, -7.66828e-02,  4.74467e-01,  ..., -8.42913e+00, -4.49072e+00, -5.53228e+00],\n",
       "            [ 4.92791e-01,  4.56345e-01,  6.32063e-01,  ..., -8.53614e+00, -4.50093e+00, -5.56321e+00],\n",
       "            [ 2.16778e-01,  1.13186e+00,  3.51918e-01,  ..., -6.99870e+00, -2.92904e+00, -5.72070e+00]],\n",
       " \n",
       "           [[-3.99943e-01, -1.50805e-01,  1.02434e-01,  ..., -7.03509e+00, -4.12126e+00, -6.51518e+00],\n",
       "            [-2.13518e-01,  3.65415e-01,  4.86163e-01,  ..., -8.13083e+00, -5.07836e+00, -6.54720e+00],\n",
       "            [-3.76944e-01,  3.94544e-01,  5.35511e-01,  ..., -8.41123e+00, -4.72656e+00, -6.71948e+00],\n",
       "            ...,\n",
       "            [-7.45143e-01,  2.53471e-01,  2.24746e-01,  ..., -8.04412e+00, -5.34447e+00, -6.76420e+00],\n",
       "            [ 3.93767e-01,  2.68533e-01,  2.98607e-01,  ..., -7.17854e+00, -4.90786e+00, -6.43760e+00],\n",
       "            [ 2.81493e-01,  1.47966e-01,  2.55903e-01,  ..., -6.56448e+00, -3.77200e+00, -6.75509e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.30895e-01, -3.14002e-01, -3.18631e-01,  ..., -5.76268e+00, -3.42350e+00, -2.77420e+00],\n",
       "            [ 4.99592e-01, -3.12883e-01, -2.62227e-01,  ..., -5.93509e+00, -3.94269e+00, -3.81211e+00],\n",
       "            [-5.76424e-01, -9.09525e-02, -1.38099e-01,  ..., -5.79687e+00, -3.59493e+00, -3.71561e+00],\n",
       "            ...,\n",
       "            [-2.28827e-01, -6.43221e-01,  3.26950e-01,  ..., -4.98463e+00, -2.75525e+00, -4.21883e+00],\n",
       "            [ 6.11037e-01, -3.58518e-01,  4.14170e-01,  ..., -5.23873e+00, -3.29590e+00, -4.27131e+00],\n",
       "            [-3.00326e-01, -4.77347e-02,  2.21175e-01,  ..., -5.65854e+00, -2.91084e+00, -3.83972e+00]],\n",
       " \n",
       "           [[-4.27582e-01, -3.41962e-01, -4.97390e-01,  ..., -5.85205e+00, -3.26168e+00, -4.71882e+00],\n",
       "            [ 3.44039e-01, -3.89045e-02, -2.64058e-01,  ..., -5.75621e+00, -3.48361e+00, -6.27103e+00],\n",
       "            [-2.52658e-01,  1.41818e-01, -6.12250e-03,  ..., -5.39300e+00, -3.09213e+00, -5.74485e+00],\n",
       "            ...,\n",
       "            [ 1.56991e-01,  4.37843e-01,  1.06943e-01,  ..., -4.64037e+00, -3.48206e+00, -5.29160e+00],\n",
       "            [ 4.00178e-01, -5.46526e-01,  4.57469e-01,  ..., -4.50220e+00, -3.51965e+00, -4.95894e+00],\n",
       "            [-4.78860e-02, -7.01338e-01,  6.29652e-02,  ..., -5.44570e+00, -3.53450e+00, -4.46340e+00]],\n",
       " \n",
       "           [[-2.28520e-01, -5.69279e-02, -2.79293e-01,  ..., -5.59656e+00, -3.42783e+00, -4.23034e+00],\n",
       "            [-5.29355e-01, -7.97471e-02, -2.19973e-01,  ..., -6.08917e+00, -3.63503e+00, -4.67580e+00],\n",
       "            [-7.59745e-01,  4.54272e-01, -5.07460e-02,  ..., -5.53215e+00, -3.19669e+00, -5.27194e+00],\n",
       "            ...,\n",
       "            [-5.50752e-01,  2.22478e-01,  7.30399e-02,  ..., -5.06367e+00, -3.54750e+00, -6.29253e+00],\n",
       "            [ 5.40988e-01,  6.80088e-01, -5.50948e-02,  ..., -5.08263e+00, -4.07677e+00, -5.25872e+00],\n",
       "            [ 8.94073e-02,  4.17884e-01, -5.46135e-02,  ..., -4.99922e+00, -3.13229e+00, -5.01732e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.53970e-01,  4.73661e-01, -4.80040e-01,  ..., -5.30747e+00, -3.36558e+00, -4.84308e+00],\n",
       "            [ 3.74809e-01, -7.31314e-01, -2.03628e-01,  ..., -5.50024e+00, -3.41124e+00, -6.56960e+00],\n",
       "            [ 2.56427e-01, -5.86388e-01, -1.31264e-01,  ..., -5.39759e+00, -2.77422e+00, -5.62789e+00],\n",
       "            ...,\n",
       "            [ 5.42729e-01,  3.64095e-01, -4.24180e-02,  ..., -4.50781e+00, -2.66650e+00, -6.84471e+00],\n",
       "            [ 1.50947e-01, -3.47583e-01, -3.24088e-01,  ..., -5.27769e+00, -3.11521e+00, -5.70481e+00],\n",
       "            [ 5.75253e-01, -4.19818e-01, -3.67558e-01,  ..., -5.40348e+00, -2.66020e+00, -4.93534e+00]],\n",
       " \n",
       "           [[-3.89291e-01,  1.34598e-02, -4.23276e-01,  ..., -5.57856e+00, -4.09076e+00, -3.85158e+00],\n",
       "            [-8.59618e-02,  9.85389e-01,  1.36368e-01,  ..., -4.93878e+00, -3.66720e+00, -5.64164e+00],\n",
       "            [ 4.94775e-01, -2.20890e-01,  4.31243e-02,  ..., -5.06161e+00, -3.20092e+00, -5.51906e+00],\n",
       "            ...,\n",
       "            [ 2.15150e-01, -1.27222e-01,  7.16006e-02,  ..., -4.73506e+00, -3.63360e+00, -6.08173e+00],\n",
       "            [ 5.65038e-01,  2.03316e-01,  2.01493e-01,  ..., -4.89381e+00, -3.52326e+00, -5.86658e+00],\n",
       "            [ 1.06114e-01,  8.58824e-01, -1.53642e-01,  ..., -5.11933e+00, -2.77881e+00, -5.09755e+00]],\n",
       " \n",
       "           [[-1.60892e-01, -1.43745e-01, -2.21753e-01,  ..., -5.66244e+00, -3.88206e+00, -2.50011e+00],\n",
       "            [-1.55377e-01,  2.85402e-01,  1.81499e-01,  ..., -5.57347e+00, -3.81400e+00, -3.73080e+00],\n",
       "            [-4.33488e-01,  3.23258e-01,  1.20306e-01,  ..., -5.57002e+00, -3.37539e+00, -3.89020e+00],\n",
       "            ...,\n",
       "            [-8.45312e-01,  1.89255e-01, -1.12165e-01,  ..., -5.77472e+00, -4.08861e+00, -3.23845e+00],\n",
       "            [ 4.93840e-01,  2.03504e-01,  2.99265e-02,  ..., -5.52032e+00, -3.83884e+00, -2.58026e+00],\n",
       "            [ 2.15678e-01,  1.31606e-01, -3.00340e-01,  ..., -6.03046e+00, -3.43922e+00, -1.57813e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.24326e-01, -1.16319e-01, -1.32213e-01,  ..., -5.13135e+00, -3.07056e+00, -2.73129e+00],\n",
       "            [ 5.06019e-01, -1.84236e-01, -5.28485e-02,  ..., -6.80514e+00, -4.55609e+00, -4.08263e+00],\n",
       "            [-5.47228e-01,  8.73599e-03,  8.89984e-02,  ..., -7.24520e+00, -4.44975e+00, -5.47429e+00],\n",
       "            ...,\n",
       "            [-2.12766e-01, -4.85121e-01,  6.91016e-01,  ..., -5.09520e+00, -2.44001e+00, -4.46152e+00],\n",
       "            [ 6.08719e-01, -2.21619e-01,  7.54281e-01,  ..., -5.38973e+00, -2.93810e+00, -4.25883e+00],\n",
       "            [-3.06880e-01,  5.02797e-02,  5.17619e-01,  ..., -4.83553e+00, -1.44253e+00, -3.84422e+00]],\n",
       " \n",
       "           [[-4.13733e-01, -2.63920e-01, -3.62135e-01,  ..., -7.43872e+00, -3.60138e+00, -4.90606e+00],\n",
       "            [ 3.57140e-01, -1.68560e-03, -1.01405e-01,  ..., -8.95686e+00, -4.38028e+00, -7.37426e+00],\n",
       "            [-2.27789e-01,  1.53752e-01,  2.28655e-01,  ..., -8.31370e+00, -3.31868e+00, -8.11526e+00],\n",
       "            ...,\n",
       "            [ 1.74644e-01,  5.04292e-01,  4.24975e-01,  ..., -6.56183e+00, -4.06974e+00, -6.19889e+00],\n",
       "            [ 4.20973e-01, -4.88104e-01,  8.92735e-01,  ..., -5.81127e+00, -3.32363e+00, -5.37963e+00],\n",
       "            [-5.54242e-02, -6.33278e-01,  3.76584e-01,  ..., -5.01481e+00, -2.49453e+00, -3.88992e+00]],\n",
       " \n",
       "           [[-2.11131e-01, -9.14961e-03, -1.30293e-01,  ..., -6.79731e+00, -3.22024e+00, -5.34381e+00],\n",
       "            [-4.93265e-01, -8.94114e-02, -5.30550e-02,  ..., -9.57865e+00, -4.48682e+00, -7.19360e+00],\n",
       "            [-7.21768e-01,  4.85903e-01,  1.60701e-01,  ..., -8.42689e+00, -3.58916e+00, -8.15619e+00],\n",
       "            ...,\n",
       "            [-5.20162e-01,  2.75395e-01,  3.32027e-01,  ..., -7.45322e+00, -3.70341e+00, -7.88624e+00],\n",
       "            [ 5.60858e-01,  7.64654e-01,  1.78235e-01,  ..., -7.16695e+00, -4.88185e+00, -6.19438e+00],\n",
       "            [ 8.98790e-02,  5.36871e-01,  2.37092e-01,  ..., -5.56551e+00, -2.16118e+00, -4.97448e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.25712e-01,  5.42617e-01, -3.45687e-01,  ..., -6.59662e+00, -2.86136e+00, -4.89789e+00],\n",
       "            [ 4.06318e-01, -7.43063e-01, -1.52636e-03,  ..., -8.90245e+00, -4.53275e+00, -6.25814e+00],\n",
       "            [ 2.79353e-01, -6.38260e-01,  1.03065e-01,  ..., -9.00309e+00, -3.99873e+00, -7.63066e+00],\n",
       "            ...,\n",
       "            [ 5.58544e-01,  3.68272e-01,  2.23305e-01,  ..., -7.86560e+00, -3.53660e+00, -8.83331e+00],\n",
       "            [ 1.62178e-01, -3.52026e-01, -1.22849e-01,  ..., -8.26271e+00, -4.06572e+00, -6.94513e+00],\n",
       "            [ 5.68987e-01, -3.76565e-01, -1.06528e-01,  ..., -6.56246e+00, -1.55580e+00, -5.65817e+00]],\n",
       " \n",
       "           [[-3.71221e-01,  8.10209e-02, -2.93325e-01,  ..., -6.31885e+00, -4.31942e+00, -3.22653e+00],\n",
       "            [-6.02497e-02,  1.06846e+00,  3.89072e-01,  ..., -7.73721e+00, -4.52161e+00, -5.40245e+00],\n",
       "            [ 5.14087e-01, -1.89792e-01,  2.89702e-01,  ..., -7.73104e+00, -3.74206e+00, -6.63431e+00],\n",
       "            ...,\n",
       "            [ 2.24163e-01, -5.09324e-02,  3.20584e-01,  ..., -7.04979e+00, -3.87495e+00, -6.09510e+00],\n",
       "            [ 5.64599e-01,  2.87878e-01,  5.42578e-01,  ..., -6.94044e+00, -3.26383e+00, -6.14509e+00],\n",
       "            [ 9.90554e-02,  1.00811e+00,  9.18045e-02,  ..., -6.11088e+00, -2.02930e+00, -5.13182e+00]],\n",
       " \n",
       "           [[-1.52642e-01, -1.59455e-01, -5.89212e-02,  ..., -5.35988e+00, -3.47328e+00, -1.47758e+00],\n",
       "            [-1.36362e-01,  2.81647e-01,  4.98713e-01,  ..., -6.03745e+00, -4.07231e+00, -2.96269e+00],\n",
       "            [-4.16367e-01,  2.89607e-01,  3.93726e-01,  ..., -6.32313e+00, -3.64518e+00, -3.37492e+00],\n",
       "            ...,\n",
       "            [-8.32811e-01,  1.92566e-01,  1.03996e-01,  ..., -5.85506e+00, -4.65077e+00, -2.21748e+00],\n",
       "            [ 4.95955e-01,  1.89769e-01,  3.26888e-01,  ..., -5.06291e+00, -3.75210e+00, -2.41075e+00],\n",
       "            [ 2.06771e-01,  1.00494e-01, -9.13715e-02,  ..., -4.75055e+00, -2.78517e+00, -2.27177e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-4.29833e-01,  8.97075e-03, -7.00593e-01,  ..., -5.51560e+00, -3.78156e+00, -1.97721e+00],\n",
       "            [-5.15759e-01,  1.10693e-01, -3.59073e-01,  ..., -5.84924e+00, -6.60562e+00, -2.84064e+00],\n",
       "            [ 1.30051e-02,  3.23382e-01, -3.77217e-01,  ..., -5.80964e+00, -6.41548e+00, -2.57404e+00],\n",
       "            ...,\n",
       "            [-4.01569e-02, -6.01021e-01, -2.76375e-01,  ..., -5.43559e+00, -5.19327e+00, -3.46552e+00],\n",
       "            [-1.56584e-01, -1.38154e-01, -1.45465e-01,  ..., -5.70913e+00, -5.47109e+00, -3.22154e+00],\n",
       "            [-9.67372e-02, -2.60339e-01, -3.63649e-02,  ..., -5.46120e+00, -3.40565e+00, -4.44988e+00]],\n",
       " \n",
       "           [[ 1.43315e-01, -7.36908e-01, -6.47614e-01,  ..., -5.90852e+00, -5.73866e+00, -2.46225e+00],\n",
       "            [ 1.38702e-01,  2.16975e-01, -5.45587e-01,  ..., -6.48279e+00, -7.98053e+00, -3.02813e+00],\n",
       "            [-7.29525e-01,  5.03627e-02, -5.50504e-01,  ..., -6.38635e+00, -7.87050e+00, -3.35101e+00],\n",
       "            ...,\n",
       "            [ 5.95174e-01, -7.51390e-01, -8.37011e-01,  ..., -6.42929e+00, -7.79003e+00, -3.15218e+00],\n",
       "            [ 4.11269e-02, -7.24811e-02, -7.36613e-01,  ..., -6.37192e+00, -7.70952e+00, -3.07984e+00],\n",
       "            [ 4.65757e-01, -5.33521e-01, -4.72125e-01,  ..., -5.47465e+00, -3.26746e+00, -3.71978e+00]],\n",
       " \n",
       "           [[-7.89046e-01,  1.61521e-01, -7.88086e-01,  ..., -5.57321e+00, -4.73142e+00, -2.28794e+00],\n",
       "            [-2.56345e-01, -1.12673e-01, -6.48704e-01,  ..., -6.39868e+00, -7.23740e+00, -2.89601e+00],\n",
       "            [-1.85426e-01,  6.39346e-01, -3.55666e-01,  ..., -6.59975e+00, -8.49343e+00, -3.51300e+00],\n",
       "            ...,\n",
       "            [-4.79496e-01, -4.31244e-02, -4.20987e-01,  ..., -6.06749e+00, -6.50207e+00, -3.28540e+00],\n",
       "            [-5.66001e-01,  7.36916e-01, -3.50991e-01,  ..., -6.22534e+00, -7.59142e+00, -4.08512e+00],\n",
       "            [ 3.50935e-01, -2.57750e-01, -6.35793e-01,  ..., -5.35690e+00, -3.64135e+00, -3.31146e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-7.60599e-01,  4.93071e-01, -8.61825e-01,  ..., -5.65787e+00, -4.65060e+00, -2.82896e+00],\n",
       "            [ 5.47404e-02, -9.12677e-01, -6.15023e-01,  ..., -6.50528e+00, -6.92909e+00, -3.48967e+00],\n",
       "            [ 1.39842e-01, -4.21889e-01, -5.65670e-01,  ..., -6.42660e+00, -6.94296e+00, -3.76629e+00],\n",
       "            ...,\n",
       "            [-9.20590e-01, -8.35120e-01,  1.46379e-02,  ..., -5.83985e+00, -6.76302e+00, -3.55612e+00],\n",
       "            [ 3.82559e-03,  4.73806e-01, -8.99710e-02,  ..., -5.85131e+00, -6.59092e+00, -4.55232e+00],\n",
       "            [ 5.17097e-01,  5.61078e-01, -4.24424e-01,  ..., -5.27903e+00, -3.29218e+00, -3.94830e+00]],\n",
       " \n",
       "           [[-8.07385e-01,  4.07253e-01, -7.56990e-01,  ..., -5.49533e+00, -3.92142e+00, -2.91286e+00],\n",
       "            [-2.45936e-01,  9.74276e-02, -5.43409e-01,  ..., -5.79475e+00, -6.18711e+00, -3.64238e+00],\n",
       "            [ 1.27806e-01, -5.32898e-01, -2.06845e-01,  ..., -5.69372e+00, -5.86949e+00, -4.13565e+00],\n",
       "            ...,\n",
       "            [-2.43814e-01,  6.76895e-01,  1.34524e-01,  ..., -5.18100e+00, -5.90988e+00, -4.20096e+00],\n",
       "            [ 2.32012e-01, -4.99870e-01, -5.92190e-02,  ..., -4.95098e+00, -5.03649e+00, -4.46533e+00],\n",
       "            [-2.17047e-01,  2.39220e-01, -3.06218e-01,  ..., -4.87780e+00, -3.99399e+00, -3.70723e+00]],\n",
       " \n",
       "           [[-2.10950e-01,  1.99337e-01, -3.51001e-01,  ..., -5.02016e+00, -3.56224e+00, -1.80722e+00],\n",
       "            [-7.28778e-01,  2.37656e-01, -2.20461e-01,  ..., -5.11742e+00, -5.36824e+00, -2.21159e+00],\n",
       "            [-1.05395e+00,  4.17182e-01, -1.78454e-01,  ..., -4.86894e+00, -5.15164e+00, -2.51691e+00],\n",
       "            ...,\n",
       "            [-3.60805e-01,  5.09714e-01,  2.47912e-01,  ..., -4.68670e+00, -5.09549e+00, -2.14229e+00],\n",
       "            [ 9.23632e-02,  4.45296e-01,  4.01948e-02,  ..., -4.64564e+00, -4.57427e+00, -2.38247e+00],\n",
       "            [-1.79351e-01,  2.58138e-01,  1.48645e-01,  ..., -4.63237e+00, -3.77766e+00, -3.51647e+00]]],\n",
       " \n",
       " \n",
       "          [[[-4.41667e-01,  3.39086e-02, -6.79214e-01,  ..., -5.30309e+00, -5.56141e+00, -6.85860e+00],\n",
       "            [-5.08945e-01,  1.54110e-01, -3.81241e-01,  ..., -7.51377e+00, -7.78338e+00, -7.44661e+00],\n",
       "            [ 2.78953e-02,  3.68492e-01, -4.01277e-01,  ..., -7.78512e+00, -7.46080e+00, -7.42215e+00],\n",
       "            ...,\n",
       "            [-2.87875e-02, -5.64406e-01, -2.48556e-01,  ..., -6.23880e+00, -5.98728e+00, -7.92213e+00],\n",
       "            [-1.44964e-01, -1.01536e-01, -7.35921e-02,  ..., -6.63972e+00, -6.94503e+00, -7.77740e+00],\n",
       "            [-8.58253e-02, -2.44018e-01,  1.23852e-01,  ..., -4.43214e+00, -4.29126e+00, -6.95162e+00]],\n",
       " \n",
       "           [[ 1.41342e-01, -7.59789e-01, -6.12276e-01,  ..., -7.67435e+00, -6.92914e+00, -7.35971e+00],\n",
       "            [ 1.43275e-01,  2.54841e-01, -6.30979e-01,  ..., -1.06879e+01, -9.42422e+00, -8.37160e+00],\n",
       "            [-7.18778e-01,  7.11574e-02, -6.39724e-01,  ..., -1.14620e+01, -8.78922e+00, -8.42599e+00],\n",
       "            ...,\n",
       "            [ 6.05645e-01, -7.88658e-01, -8.95365e-01,  ..., -1.12603e+01, -9.04565e+00, -8.89822e+00],\n",
       "            [ 4.92678e-02, -5.39316e-02, -7.75714e-01,  ..., -1.04682e+01, -9.10571e+00, -8.78779e+00],\n",
       "            [ 4.74364e-01, -5.24933e-01, -3.89985e-01,  ..., -4.70388e+00, -4.38190e+00, -7.13287e+00]],\n",
       " \n",
       "           [[-7.99999e-01,  1.54068e-01, -8.17426e-01,  ..., -6.64489e+00, -5.61523e+00, -7.29088e+00],\n",
       "            [-2.58020e-01, -9.16972e-02, -7.76953e-01,  ..., -1.08916e+01, -8.50033e+00, -8.06605e+00],\n",
       "            [-1.72322e-01,  6.64652e-01, -4.79242e-01,  ..., -1.19906e+01, -9.06217e+00, -8.02392e+00],\n",
       "            ...,\n",
       "            [-4.78978e-01, -3.12751e-02, -5.08278e-01,  ..., -1.15642e+01, -8.37041e+00, -7.89233e+00],\n",
       "            [-5.51619e-01,  7.73110e-01, -4.05052e-01,  ..., -1.13670e+01, -8.85517e+00, -8.67973e+00],\n",
       "            [ 3.64118e-01, -2.63650e-01, -6.30291e-01,  ..., -6.81863e+00, -5.20071e+00, -7.50179e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-7.66207e-01,  5.13449e-01, -8.43796e-01,  ..., -6.36754e+00, -5.17877e+00, -7.14491e+00],\n",
       "            [ 5.08093e-02, -9.25165e-01, -6.62987e-01,  ..., -1.09031e+01, -7.97981e+00, -8.09995e+00],\n",
       "            [ 1.40935e-01, -3.82814e-01, -6.38788e-01,  ..., -1.19445e+01, -8.51332e+00, -8.50211e+00],\n",
       "            ...,\n",
       "            [-9.18488e-01, -8.03540e-01, -2.78023e-02,  ..., -1.01930e+01, -8.38815e+00, -7.96385e+00],\n",
       "            [ 1.38308e-02,  5.31270e-01, -9.76200e-02,  ..., -1.05281e+01, -8.59629e+00, -8.21187e+00],\n",
       "            [ 5.24570e-01,  5.94851e-01, -3.95646e-01,  ..., -7.13774e+00, -5.40785e+00, -7.36020e+00]],\n",
       " \n",
       "           [[-8.21446e-01,  4.15421e-01, -7.27061e-01,  ..., -6.33274e+00, -5.00001e+00, -6.84638e+00],\n",
       "            [-2.50477e-01,  1.11943e-01, -5.92794e-01,  ..., -1.02428e+01, -7.75044e+00, -7.57308e+00],\n",
       "            [ 1.26928e-01, -4.92987e-01, -2.47735e-01,  ..., -1.08920e+01, -7.56747e+00, -7.59248e+00],\n",
       "            ...,\n",
       "            [-2.36055e-01,  6.97096e-01,  2.33608e-01,  ..., -8.61952e+00, -6.93042e+00, -7.39166e+00],\n",
       "            [ 2.30509e-01, -4.75467e-01,  2.27484e-02,  ..., -9.34407e+00, -7.04778e+00, -7.27631e+00],\n",
       "            [-2.02043e-01,  2.43080e-01, -2.23802e-01,  ..., -6.47614e+00, -5.32735e+00, -6.86440e+00]],\n",
       " \n",
       "           [[-2.20011e-01,  1.93664e-01, -3.02791e-01,  ..., -4.89906e+00, -4.69708e+00, -5.21612e+00],\n",
       "            [-7.21457e-01,  2.40416e-01, -2.18120e-01,  ..., -5.98045e+00, -5.73313e+00, -6.06184e+00],\n",
       "            [-1.04546e+00,  4.18478e-01, -1.74550e-01,  ..., -6.00448e+00, -5.37216e+00, -6.46112e+00],\n",
       "            ...,\n",
       "            [-3.54405e-01,  4.88294e-01,  3.41744e-01,  ..., -5.11192e+00, -5.44782e+00, -5.70084e+00],\n",
       "            [ 1.00019e-01,  4.07031e-01,  1.31933e-01,  ..., -4.50444e+00, -4.96368e+00, -5.66432e+00],\n",
       "            [-1.61727e-01,  2.22381e-01,  3.48966e-01,  ..., -3.43624e+00, -4.14836e+00, -5.65704e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.90008e-01, -1.02519e-01, -7.33007e-01,  ..., -4.52173e+00, -2.87153e+00, -4.71236e+00],\n",
       "            [-4.60667e-01,  6.83446e-02, -4.14324e-01,  ..., -6.32112e+00, -5.03290e+00, -5.81848e+00],\n",
       "            [ 4.97821e-02,  2.75112e-01, -3.95566e-01,  ..., -6.43497e+00, -4.66392e+00, -6.27954e+00],\n",
       "            ...,\n",
       "            [-1.95171e-02, -6.05852e-01, -3.68768e-01,  ..., -5.03915e+00, -3.29116e+00, -6.30870e+00],\n",
       "            [-1.18655e-01, -1.91238e-01, -2.30763e-01,  ..., -5.41298e+00, -3.40862e+00, -5.84083e+00],\n",
       "            [-1.42251e-01, -2.76425e-01, -3.66964e-01,  ..., -4.60284e+00, -2.41398e+00, -4.83832e+00]],\n",
       " \n",
       "           [[ 2.52677e-01, -8.74280e-01, -7.60429e-01,  ..., -5.74226e+00, -4.32330e+00, -5.18990e+00],\n",
       "            [ 2.00642e-01,  1.94261e-01, -4.99352e-01,  ..., -7.54110e+00, -5.27997e+00, -6.88312e+00],\n",
       "            [-6.73729e-01,  5.21050e-02, -5.71607e-01,  ..., -7.96360e+00, -5.31039e+00, -7.13132e+00],\n",
       "            ...,\n",
       "            [ 6.05859e-01, -9.01395e-01, -8.13429e-01,  ..., -7.37111e+00, -4.87072e+00, -7.65049e+00],\n",
       "            [ 7.95756e-02, -1.15634e-01, -7.44031e-01,  ..., -6.97683e+00, -5.09136e+00, -7.02913e+00],\n",
       "            [ 4.04350e-01, -5.88561e-01, -6.57584e-01,  ..., -4.28998e+00, -2.42759e+00, -4.90278e+00]],\n",
       " \n",
       "           [[-6.22612e-01, -3.73116e-02, -7.43112e-01,  ..., -5.41749e+00, -3.00970e+00, -5.49646e+00],\n",
       "            [-1.90680e-01, -2.33395e-01, -6.37258e-01,  ..., -7.90372e+00, -4.72632e+00, -7.03409e+00],\n",
       "            [-1.39236e-01,  5.57370e-01, -3.41091e-01,  ..., -9.49230e+00, -5.44448e+00, -7.58856e+00],\n",
       "            ...,\n",
       "            [-4.34130e-01, -1.24196e-01, -5.29563e-01,  ..., -8.22915e+00, -3.77166e+00, -7.67746e+00],\n",
       "            [-5.53782e-01,  6.60720e-01, -4.06347e-01,  ..., -8.31594e+00, -4.68659e+00, -8.04870e+00],\n",
       "            [ 2.80197e-01, -3.48348e-01, -6.86228e-01,  ..., -5.61373e+00, -2.26417e+00, -5.96706e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.92157e-01,  4.07766e-01, -8.46440e-01,  ..., -5.42556e+00, -3.41951e+00, -4.84089e+00],\n",
       "            [ 1.21167e-01, -9.54085e-01, -7.11412e-01,  ..., -7.95653e+00, -5.12357e+00, -6.54877e+00],\n",
       "            [ 2.12050e-01, -4.47377e-01, -5.07816e-01,  ..., -8.02043e+00, -4.50119e+00, -7.40893e+00],\n",
       "            ...,\n",
       "            [-8.69429e-01, -8.48577e-01, -2.09828e-01,  ..., -7.66430e+00, -4.58608e+00, -6.56250e+00],\n",
       "            [ 4.62833e-02,  4.90129e-01, -1.96678e-01,  ..., -7.45755e+00, -4.15590e+00, -6.60672e+00],\n",
       "            [ 4.44013e-01,  5.57395e-01, -5.73602e-01,  ..., -5.66164e+00, -2.54505e+00, -5.74494e+00]],\n",
       " \n",
       "           [[-6.42307e-01,  3.13853e-01, -7.99539e-01,  ..., -5.17348e+00, -3.06348e+00, -4.58147e+00],\n",
       "            [-1.47769e-01, -8.84368e-03, -5.45518e-01,  ..., -7.03949e+00, -4.46828e+00, -6.26634e+00],\n",
       "            [ 2.15788e-01, -6.21677e-01, -1.88581e-01,  ..., -7.70584e+00, -3.88153e+00, -6.81945e+00],\n",
       "            ...,\n",
       "            [-1.85895e-01,  6.51100e-01, -8.63826e-03,  ..., -6.76866e+00, -4.01853e+00, -6.40973e+00],\n",
       "            [ 2.82131e-01, -5.40879e-01, -1.52346e-01,  ..., -6.84012e+00, -3.20517e+00, -6.88473e+00],\n",
       "            [-2.92279e-01,  2.01620e-01, -4.14449e-01,  ..., -5.31164e+00, -2.94316e+00, -5.63831e+00]],\n",
       " \n",
       "           [[-1.15531e-01,  1.28058e-01, -6.01551e-01,  ..., -4.94448e+00, -3.61813e+00, -2.99508e+00],\n",
       "            [-6.80582e-01,  1.77212e-01, -3.66914e-01,  ..., -5.43208e+00, -4.53518e+00, -4.19648e+00],\n",
       "            [-1.00190e+00,  3.26540e-01, -2.92972e-01,  ..., -5.21744e+00, -4.13462e+00, -4.84750e+00],\n",
       "            ...,\n",
       "            [-3.15852e-01,  4.07035e-01, -4.11969e-02,  ..., -5.10846e+00, -4.31751e+00, -4.31238e+00],\n",
       "            [ 1.24720e-01,  3.32626e-01, -2.10389e-01,  ..., -4.55623e+00, -4.09629e+00, -4.18789e+00],\n",
       "            [-2.44649e-01,  1.60214e-01, -2.95506e-01,  ..., -4.25031e+00, -3.66998e+00, -3.83435e+00]]]]], device='cuda:0'))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qat_model(torch.rand((1,3,640,640)).cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f4453b9a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:07.571968 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.572519 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.574281 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.574583 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.579425 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.579825 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.581447 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.581789 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.586458 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.586788 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.588322 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.588615 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.593417 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.593933 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.595615 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.596001 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.600591 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.600960 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.602448 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.602696 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.607386 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.607643 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.609142 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.609388 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.615207 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.615464 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.616968 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.617216 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.624042 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.624329 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.625872 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.626143 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.630810 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.631071 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.632568 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.632854 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.637466 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.637720 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.639211 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.639455 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.644086 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.644340 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.645904 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.646150 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.651725 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.651981 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.653504 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.653755 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.660322 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.660588 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.662085 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.662329 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.666913 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.667200 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.668714 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.669320 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.675355 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.675704 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.677577 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.677853 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.681951 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.682204 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.683434 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.683672 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.688115 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.688364 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.689590 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.689836 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.695068 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.695321 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.696544 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.696786 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.700462 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.700711 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:07.701926 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.702162 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.705913 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.706160 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.707381 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.707617 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.711303 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.711558 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.712790 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.713031 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.718255 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.718519 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.719746 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.719986 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.723684 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.724010 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.725236 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.725497 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.729201 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.729449 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.730672 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.730911 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.740007 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.740343 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.741530 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.741771 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.746176 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.746427 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.747645 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.747882 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.751556 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.751803 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.753027 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.753280 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.757851 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.758136 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.759375 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.759614 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.764033 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.764282 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.765503 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.765742 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.769402 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.769624 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.770771 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.770981 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.774586 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.775021 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.776160 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.776372 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.779951 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.780211 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.781405 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.781619 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.785967 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.786189 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.787340 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.787560 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.791169 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.791446 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.792583 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.792794 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.797019 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.797291 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.798391 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.798641 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.803481 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.803786 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.804968 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.805238 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.809298 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.809589 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.810706 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.810912 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:07.814574 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.814785 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.815910 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.816124 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.819688 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.819947 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.821134 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.821339 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.825651 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.825903 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.827014 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.827224 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.830750 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.830965 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.832205 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.832407 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.838043 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.838307 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.839422 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.839632 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.843305 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.843516 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.844826 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.845041 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.848651 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.848968 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.850119 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.850322 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.853871 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.854127 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.855314 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.855534 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.860117 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.860358 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.861501 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.861714 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.865293 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.865545 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.866654 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.866855 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.871578 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.871851 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.872966 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.873172 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.876690 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.876946 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.878048 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.878259 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.881784 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.881995 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.883119 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.883323 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.886953 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.887362 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.888578 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.888865 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.893492 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.893774 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.894915 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.895125 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.898757 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.899012 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.900133 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.900339 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.904376 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.904634 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.905774 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.905982 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.909569 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.909826 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.910948 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:07.911154 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.047414 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.047766 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:08.049036 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.049306 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.053107 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.053398 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.054484 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.054757 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.058460 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.058743 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.059817 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.060065 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.064140 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.064404 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.065544 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.065766 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.069321 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.069642 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.071041 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.071307 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.074781 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.075054 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.076351 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.076564 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.080842 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.081107 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.082164 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.082439 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.087941 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.088215 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.089388 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.089606 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.093084 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.093299 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.094487 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.094698 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.098455 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.098703 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.099814 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.100023 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.103539 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.103787 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.105024 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.105237 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.109908 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.110172 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.111269 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.111479 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.116504 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.116720 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.117873 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.118085 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.121987 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.122199 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.123321 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.123524 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.126998 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.127211 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.128300 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.128502 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.131879 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.132086 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.133224 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.133427 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.137558 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.137771 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.138918 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.139244 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.144458 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.144696 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.145864 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.146069 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.149611 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.149823 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.150915 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.151173 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:08.155211 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.155464 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.156790 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.157036 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.160909 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.161161 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.162359 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.162604 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.166823 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.167076 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.168267 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.168509 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.172037 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.172286 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.173581 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.173822 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.177310 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.177557 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.178749 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.178990 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.187588 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.187903 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.189184 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.189496 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.193731 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.194069 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.195254 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.195500 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.199162 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.199417 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.200636 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.200952 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.205070 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.205332 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.206624 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.206936 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:08.211385 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.211788 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.213011 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.213257 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.216838 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.217093 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.218309 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.218553 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.222132 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.222503 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.223678 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.223989 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.227489 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.227741 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.228960 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.229206 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.233451 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.233705 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.234917 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.235230 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.238784 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.239104 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.240310 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.240558 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.245018 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.245276 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.246492 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.246805 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.251185 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.251462 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.252718 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.253107 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.256595 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.256921 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.258198 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.258476 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.262071 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.262360 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.263690 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.263957 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.268724 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.269034 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.270103 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.270362 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.275418 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.275722 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.277867 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.278146 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.281611 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.281895 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.282936 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.283211 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.287290 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.287561 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.288610 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.288926 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.292341 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.292622 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.294292 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.294506 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.297953 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.298170 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.299422 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.299684 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.303199 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.303462 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.304622 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.304847 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.309233 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.309517 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.310609 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.310825 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.314304 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.314557 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0615 14:21:08.315640 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.315844 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.320095 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.320371 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.321533 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.321779 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.325308 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.325578 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.326616 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.326879 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.330248 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.330585 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.331622 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.331896 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.335284 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.335565 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.336597 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.336861 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.341308 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.341596 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.342635 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.342892 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.346330 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.346616 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.347646 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.347905 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.351299 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.351576 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.353597 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.353812 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.357416 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.357718 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.358852 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n",
      "E0615 14:21:08.359116 140418879964992 tensor_quantizer.py:120] Fake quantize mode doesn't use scale explicitly!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n",
      "fuseforward training\n"
     ]
    }
   ],
   "source": [
    "quant_nn.TensorQuantizer.use_fb_fake_quant = True\n",
    "with torch.no_grad():\n",
    "    data = iter(dataloader)\n",
    "    images, _,_,_ = next(data)\n",
    "    images = (images/255.0).cuda()\n",
    "    jit_model = torch.jit.trace(qat_model, images.to(\"cuda\"))\n",
    "    torch.jit.save(jit_model, \"trained_vgg16_qat.jit.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c3814d11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (model): Sequential(\n",
       "    (0): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0000 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.8467, 9.8950](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (1): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=21.9565 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1028, 1.0065](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (2): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=31.6012 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.3582, 1.2886](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (3): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=31.6012 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.3486, 1.3525](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (4): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.4433 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1656, 0.7424](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (5): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=17.1036 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1998, 0.4605](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (6): Concat()\n",
       "    (7): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=17.1036 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.2378, 0.6717](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (8): MP(\n",
       "      (m): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (9): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.7819 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1816, 0.7992](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (10): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.7819 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1792, 0.7680](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (11): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.7092 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1380, 0.3033](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (12): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.3041 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1401, 0.3056](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (13): Concat()\n",
       "    (14): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=13.7202 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1422, 0.5593](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (15): MP(\n",
       "      (m): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (16): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=12.3185 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1325, 0.5530](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (17): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=12.3185 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1662, 0.5371](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (18): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.5341 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0773, 0.2254](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (19): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.2964 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0871, 0.2030](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (20): Concat()\n",
       "    (21): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.4538 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1095, 0.7497](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (22): MP(\n",
       "      (m): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (23): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.7248 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0860, 0.2745](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (24): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.7248 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1073, 0.5097](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (25): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.7734 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0468, 0.1695](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (26): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.2573 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0491, 0.1701](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (27): Concat()\n",
       "    (28): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        1024, 512, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.3556 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0789, 0.1798](512) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (29): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=6.8431 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1131, 0.2692](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (30): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=6.8431 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0836, 0.2447](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (31): SP(\n",
       "      (m): MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (32): SP(\n",
       "      (m): MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (33): SP(\n",
       "      (m): MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (34): Concat()\n",
       "    (35): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=6.9201 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0546, 0.1576](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (36): Concat()\n",
       "    (37): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.8761 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1058, 0.3790](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (38): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.9318 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1772, 0.4712](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (39): Upsample(scale_factor=2.0, mode=nearest)\n",
       "    (40): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.7248 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1236, 0.3978](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (41): Concat()\n",
       "    (42): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.6507 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1526, 0.4597](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (43): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.6507 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1507, 0.3301](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (44): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=7.6914 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0871, 0.3623](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (45): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=7.5127 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0984, 0.2614](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (46): Concat()\n",
       "    (47): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.4010 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1394, 0.5261](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (48): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.0094 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.2262, 0.8891](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (49): Upsample(scale_factor=2.0, mode=nearest)\n",
       "    (50): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=12.3185 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1668, 0.8992](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (51): Concat()\n",
       "    (52): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.8600 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.2115, 0.6500](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (53): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.8600 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1715, 0.4547](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (54): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=7.2781 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.2098, 0.6950](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (55): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.7555 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.2554, 0.8387](32) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (56): Concat()\n",
       "    (57): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.4869 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1830, 0.7702](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (58): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.5577 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1000, 0.3975](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (59): Concat()\n",
       "    (60): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.9869 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1485, 0.5582](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (61): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.9869 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1336, 0.6178](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (62): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.7306 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1078, 0.5156](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (63): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=9.3281 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1375, 0.6402](64) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (64): Concat()\n",
       "    (65): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.4749 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1214, 0.9532](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (66): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.9742 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0747, 0.3336](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (67): Concat()\n",
       "    (68): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.2716 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0944, 0.3416](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (69): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=10.2716 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0846, 0.4276](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (70): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=12.6105 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0681, 0.4963](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (71): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.0984 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0713, 0.5777](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (72): Concat()\n",
       "    (73): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=12.6105 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0650, 0.7201](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (74): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.5577 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1087, 1.2581](128) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (75): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=11.9742 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0504, 1.1487](256) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (76): Conv(\n",
       "      (conv): QuantConv2d(\n",
       "        256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=8.8608 calibrator=MaxCalibrator scale=1.0 quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0312, 0.6864](512) calibrator=MaxCalibrator scale=1.0 quant)\n",
       "      )\n",
       "      (act): LeakyReLU(negative_slope=0.1, inplace=True)\n",
       "    )\n",
       "    (77): IDetect(\n",
       "      (m): ModuleList(\n",
       "        (0): Conv2d(128, 45, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (1): Conv2d(256, 45, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (2): Conv2d(512, 45, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (ia): ModuleList(\n",
       "        (0): ImplicitA()\n",
       "        (1): ImplicitA()\n",
       "        (2): ImplicitA()\n",
       "      )\n",
       "      (im): ModuleList(\n",
       "        (0): ImplicitM()\n",
       "        (1): ImplicitM()\n",
       "        (2): ImplicitM()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_model.cpu().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "58cf6dbc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[2.63411e+00, 5.58906e+00, 4.83296e+00,  ..., 1.07195e-03, 2.33448e-03, 7.97935e-03],\n",
       "          [7.66330e+00, 4.77671e+00, 5.89252e+00,  ..., 1.31260e-03, 4.92726e-03, 1.87256e-02],\n",
       "          [1.82627e+01, 4.07512e+00, 5.74930e+00,  ..., 7.25613e-04, 2.68547e-03, 1.27729e-02],\n",
       "          ...,\n",
       "          [5.98103e+02, 3.75296e+02, 3.74372e+01,  ..., 7.77785e-03, 1.46844e-02, 8.82231e-03],\n",
       "          [6.20645e+02, 3.76648e+02, 5.15256e+01,  ..., 1.02470e-02, 3.01174e-02, 9.34405e-03],\n",
       "          [6.49215e+02, 3.71123e+02, 4.15049e+01,  ..., 1.01309e-02, 2.29706e-02, 2.25935e-02]]]),\n",
       " [tensor([[[[[-3.44849e-01,  4.02616e-01,  4.85726e-01,  ..., -6.83720e+00, -6.05763e+00, -4.82289e+00],\n",
       "             [-1.21421e+00,  1.94792e-01,  7.70379e-01,  ..., -6.63443e+00, -5.30803e+00, -3.95896e+00],\n",
       "             [-4.41349e-01,  1.87811e-02,  7.32000e-01,  ..., -7.22777e+00, -5.91721e+00, -4.34757e+00],\n",
       "             ...,\n",
       "             [-6.40948e-01, -2.28585e-02,  4.00428e-01,  ..., -8.06027e+00, -6.53119e+00, -5.94431e+00],\n",
       "             [-7.56900e-01, -1.96989e-01,  4.91586e-01,  ..., -7.46453e+00, -6.33149e+00, -5.30963e+00],\n",
       "             [-3.34654e-01, -2.48650e-01,  5.66298e-01,  ..., -5.92954e+00, -5.69576e+00, -4.20253e+00]],\n",
       "  \n",
       "            [[-1.55896e-01, -8.75497e-01,  6.05756e-01,  ..., -6.57748e+00, -6.31385e+00, -4.75126e+00],\n",
       "             [-9.78089e-01, -3.99484e-02,  7.85930e-01,  ..., -6.87777e+00, -5.84006e+00, -5.10151e+00],\n",
       "             [-7.05200e-01, -9.06101e-01,  9.06488e-01,  ..., -6.83421e+00, -5.54618e+00, -3.96156e+00],\n",
       "             ...,\n",
       "             [-8.90648e-01, -7.12583e-01,  6.14261e-01,  ..., -7.07983e+00, -5.77256e+00, -4.46671e+00],\n",
       "             [-1.89976e-01, -6.97529e-01,  6.74654e-01,  ..., -6.73983e+00, -5.48049e+00, -4.46351e+00],\n",
       "             [-1.83038e-01, -4.90745e-01,  5.99527e-01,  ..., -6.15106e+00, -5.58590e+00, -4.51628e+00]],\n",
       "  \n",
       "            [[-2.11401e-01, -1.26138e+00,  2.65478e-01,  ..., -5.65492e+00, -5.37857e+00, -3.11503e+00],\n",
       "             [-7.14134e-01, -7.87977e-01,  8.86640e-01,  ..., -6.23157e+00, -5.27310e+00, -4.18436e+00],\n",
       "             [-6.66258e-01, -2.43708e-01,  8.44101e-01,  ..., -6.81488e+00, -4.90440e+00, -3.90026e+00],\n",
       "             ...,\n",
       "             [-9.84127e-01,  7.90163e-01,  5.87699e-01,  ..., -7.77436e+00, -6.58354e+00, -5.81991e+00],\n",
       "             [ 7.38199e-01, -1.88233e-01,  4.76911e-01,  ..., -7.83041e+00, -6.30508e+00, -5.34029e+00],\n",
       "             [-4.13561e-01, -4.47150e-01,  5.43839e-01,  ..., -6.49010e+00, -5.63699e+00, -4.23927e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[-3.72952e-01, -7.13346e-02,  4.65882e-01,  ..., -5.90029e+00, -4.66541e+00, -3.42114e+00],\n",
       "             [-4.40362e-01,  2.28737e-01,  8.05636e-01,  ..., -6.21523e+00, -4.55559e+00, -3.75949e+00],\n",
       "             [-7.19300e-02,  1.10641e-01,  1.03238e+00,  ..., -6.76284e+00, -5.02092e+00, -4.22439e+00],\n",
       "             ...,\n",
       "             [-1.56504e-01, -2.02342e-01,  6.87731e-01,  ..., -6.87775e+00, -4.63905e+00, -4.38823e+00],\n",
       "             [-3.80025e-01,  6.05737e-01,  5.90191e-01,  ..., -7.13265e+00, -4.98042e+00, -4.26341e+00],\n",
       "             [ 3.32357e-01, -1.95267e-01,  4.43825e-01,  ..., -6.01366e+00, -5.33416e+00, -3.16994e+00]],\n",
       "  \n",
       "            [[-3.28202e-01,  4.11771e-01,  6.55620e-01,  ..., -5.97670e+00, -5.31802e+00, -3.38632e+00],\n",
       "             [-9.76868e-01,  6.77558e-01,  9.37257e-01,  ..., -6.49403e+00, -5.23652e+00, -4.26687e+00],\n",
       "             [-2.73159e-01,  4.79092e-01,  7.86564e-01,  ..., -6.87223e+00, -5.56919e+00, -4.24553e+00],\n",
       "             ...,\n",
       "             [-3.51933e-01,  1.02521e+00,  6.64751e-01,  ..., -7.59366e+00, -5.79815e+00, -4.51734e+00],\n",
       "             [ 3.03777e-01,  3.28511e-01,  6.35059e-01,  ..., -7.40796e+00, -5.81534e+00, -4.00374e+00],\n",
       "             [-3.81882e-01,  1.67175e-01,  6.20553e-01,  ..., -6.04484e+00, -5.00260e+00, -2.84943e+00]],\n",
       "  \n",
       "            [[-2.22810e-01, -3.75293e-01,  7.47675e-01,  ..., -6.16898e+00, -5.82824e+00, -3.06877e+00],\n",
       "             [-6.26413e-01, -7.45029e-02,  1.00365e+00,  ..., -6.35780e+00, -5.49641e+00, -3.11413e+00],\n",
       "             [-8.30782e-01,  1.89675e-02,  1.16508e+00,  ..., -6.03455e+00, -5.39316e+00, -2.48042e+00],\n",
       "             ...,\n",
       "             [-5.49232e-01, -1.95879e-01,  7.29075e-01,  ..., -7.95911e+00, -6.68135e+00, -4.10819e+00],\n",
       "             [-1.08705e+00,  3.85913e-02,  8.62040e-01,  ..., -6.39495e+00, -5.55052e+00, -2.11744e+00],\n",
       "             [-3.94292e-01, -3.28539e-01,  7.78327e-01,  ..., -5.84144e+00, -5.04039e+00, -2.20706e+00]]],\n",
       "  \n",
       "  \n",
       "           [[[ 1.97583e-01,  6.36815e-01,  3.31500e-01,  ..., -6.95642e+00, -6.27625e+00, -4.01895e+00],\n",
       "             [ 3.95722e-01,  3.65422e-01,  6.02839e-01,  ..., -6.99351e+00, -5.56026e+00, -3.27135e+00],\n",
       "             [ 3.97005e-01,  2.32952e-01,  5.91866e-01,  ..., -7.78621e+00, -5.99860e+00, -3.56318e+00],\n",
       "             ...,\n",
       "             [ 3.36960e-01,  3.52730e-02,  1.67843e-01,  ..., -8.96012e+00, -6.92887e+00, -5.51045e+00],\n",
       "             [ 3.55217e-01, -1.52279e-01,  2.26783e-01,  ..., -8.20306e+00, -6.87806e+00, -4.87305e+00],\n",
       "             [ 1.45103e-01, -2.13879e-01,  2.29469e-01,  ..., -5.77748e+00, -6.00784e+00, -3.37779e+00]],\n",
       "  \n",
       "            [[ 3.56140e-02, -5.01276e-01,  4.53498e-01,  ..., -7.11461e+00, -6.55448e+00, -4.34184e+00],\n",
       "             [ 4.37738e-02,  3.11139e-02,  6.61186e-01,  ..., -7.75022e+00, -6.30821e+00, -5.04353e+00],\n",
       "             [ 4.28816e-01, -6.94509e-01,  9.92130e-01,  ..., -7.38837e+00, -5.71114e+00, -3.91291e+00],\n",
       "             ...,\n",
       "             [ 3.90189e-01, -3.76785e-01,  5.03642e-01,  ..., -7.84660e+00, -6.06327e+00, -4.61044e+00],\n",
       "             [ 1.08113e+00, -4.37477e-01,  5.54700e-01,  ..., -7.41358e+00, -5.79630e+00, -4.48979e+00],\n",
       "             [ 3.22488e-01, -7.42892e-02,  3.26508e-01,  ..., -6.36326e+00, -5.74865e+00, -4.15429e+00]],\n",
       "  \n",
       "            [[-1.81315e-01, -9.88007e-01,  1.99065e-01,  ..., -5.92449e+00, -5.52016e+00, -2.68039e+00],\n",
       "             [ 1.60957e-02, -5.44333e-01,  1.00883e+00,  ..., -6.62353e+00, -5.46009e+00, -3.93500e+00],\n",
       "             [ 5.73824e-01, -3.54545e-02,  9.50250e-01,  ..., -7.42724e+00, -5.02757e+00, -3.88279e+00],\n",
       "             ...,\n",
       "             [-3.57793e-01,  1.14095e+00,  4.29396e-01,  ..., -8.87060e+00, -6.83049e+00, -5.53133e+00],\n",
       "             [ 1.57846e+00,  1.30918e-01,  3.13901e-01,  ..., -8.95779e+00, -6.25175e+00, -5.04300e+00],\n",
       "             [ 2.71544e-01, -7.69796e-02,  2.73758e-01,  ..., -7.03317e+00, -5.78047e+00, -3.97293e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[-1.11799e-01,  3.70543e-01,  5.08370e-01,  ..., -6.10275e+00, -4.55528e+00, -3.05007e+00],\n",
       "             [ 8.84304e-01,  4.96219e-01,  9.81157e-01,  ..., -6.76468e+00, -5.03191e+00, -3.75455e+00],\n",
       "             [ 1.19144e+00,  2.70679e-01,  1.03702e+00,  ..., -7.55624e+00, -5.16051e+00, -4.04057e+00],\n",
       "             ...,\n",
       "             [ 7.85836e-01,  1.27834e-01,  8.01338e-01,  ..., -7.63600e+00, -4.77693e+00, -4.20782e+00],\n",
       "             [ 1.03302e+00,  7.44668e-01,  6.52477e-01,  ..., -8.27267e+00, -5.21111e+00, -4.12793e+00],\n",
       "             [ 2.70460e-01,  2.58066e-01,  1.92518e-01,  ..., -6.59894e+00, -5.11846e+00, -2.75384e+00]],\n",
       "  \n",
       "            [[ 4.96723e-02,  9.58175e-01,  6.54425e-01,  ..., -6.13229e+00, -5.15977e+00, -3.02017e+00],\n",
       "             [ 2.93733e-01,  1.01213e+00,  9.65325e-01,  ..., -6.88484e+00, -5.27276e+00, -3.96665e+00],\n",
       "             [ 6.68884e-01,  8.26566e-01,  8.08294e-01,  ..., -7.67183e+00, -5.49862e+00, -4.15649e+00],\n",
       "             ...,\n",
       "             [ 3.54482e-01,  1.33375e+00,  5.85589e-01,  ..., -8.68869e+00, -5.82565e+00, -4.40354e+00],\n",
       "             [ 1.34109e+00,  6.77787e-01,  7.21367e-01,  ..., -8.26602e+00, -5.61721e+00, -3.80526e+00],\n",
       "             [ 1.72132e-01,  7.01558e-01,  4.54897e-01,  ..., -6.28629e+00, -4.77301e+00, -2.86721e+00]],\n",
       "  \n",
       "            [[ 1.14142e-01, -3.93357e-01,  5.59983e-01,  ..., -6.22430e+00, -5.74930e+00, -2.81103e+00],\n",
       "             [ 9.35906e-01, -2.77980e-01,  7.47034e-01,  ..., -6.73700e+00, -5.56862e+00, -3.18599e+00],\n",
       "             [ 4.37348e-01, -1.37190e-01,  1.06956e+00,  ..., -6.42027e+00, -5.18461e+00, -2.52704e+00],\n",
       "             ...,\n",
       "             [ 2.76006e-01, -2.68770e-01,  4.39317e-01,  ..., -9.07695e+00, -6.75565e+00, -4.24187e+00],\n",
       "             [ 4.32224e-01, -1.61619e-01,  7.54992e-01,  ..., -7.13107e+00, -5.44647e+00, -2.09924e+00],\n",
       "             [-3.41542e-02, -2.42757e-01,  5.37405e-01,  ..., -5.85822e+00, -4.80580e+00, -2.19148e+00]]],\n",
       "  \n",
       "  \n",
       "           [[[-1.38255e-01,  7.72659e-01, -2.32563e-01,  ..., -5.54700e+00, -5.73791e+00, -3.08956e+00],\n",
       "             [-4.12358e-01,  4.23481e-01,  2.17431e-01,  ..., -5.28254e+00, -4.51926e+00, -2.62165e+00],\n",
       "             [-1.52591e-01,  3.03734e-01,  1.70890e-01,  ..., -6.36027e+00, -5.13469e+00, -3.72546e+00],\n",
       "             ...,\n",
       "             [-3.12418e-01,  1.43209e-03, -2.63142e-01,  ..., -7.54928e+00, -6.11936e+00, -5.19649e+00],\n",
       "             [-3.48686e-01, -1.33119e-01, -1.95770e-01,  ..., -6.66763e+00, -6.03959e+00, -3.71699e+00],\n",
       "             [-6.71993e-02, -1.59969e-01, -1.97800e-01,  ..., -4.33888e+00, -5.74843e+00, -1.92886e+00]],\n",
       "  \n",
       "            [[-8.42384e-02, -4.37691e-02, -1.78876e-01,  ..., -6.39918e+00, -5.91602e+00, -3.81252e+00],\n",
       "             [-6.41984e-01,  2.07491e-01,  1.99007e-01,  ..., -6.80178e+00, -5.22314e+00, -4.69453e+00],\n",
       "             [-3.59936e-01, -2.42251e-01,  5.26314e-01,  ..., -6.45947e+00, -4.62907e+00, -4.27174e+00],\n",
       "             ...,\n",
       "             [-5.18922e-01,  2.42933e-01,  5.78711e-02,  ..., -6.91008e+00, -5.14909e+00, -4.33404e+00],\n",
       "             [ 2.14417e-01, -3.59575e-02,  6.17726e-02,  ..., -6.44369e+00, -4.99778e+00, -3.94079e+00],\n",
       "             [-7.97237e-02,  4.33191e-01, -1.46015e-01,  ..., -5.27878e+00, -5.23587e+00, -3.19356e+00]],\n",
       "  \n",
       "            [[-1.73088e-01, -4.19920e-01, -3.26133e-01,  ..., -5.58253e+00, -4.94078e+00, -2.24310e+00],\n",
       "             [-4.29254e-01, -1.33981e-01,  4.52436e-01,  ..., -5.91584e+00, -4.56546e+00, -4.04166e+00],\n",
       "             [-4.23307e-01,  3.47240e-01,  4.98797e-01,  ..., -6.70865e+00, -4.08447e+00, -4.58454e+00],\n",
       "             ...,\n",
       "             [-6.98706e-01,  1.40950e+00, -2.09248e-01,  ..., -7.83830e+00, -6.16296e+00, -5.00003e+00],\n",
       "             [ 1.01389e+00,  6.89107e-01, -2.19617e-01,  ..., -8.37354e+00, -5.58242e+00, -4.69794e+00],\n",
       "             [-4.44210e-02,  3.42165e-01, -2.61670e-01,  ..., -6.11048e+00, -5.23940e+00, -3.65743e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[-8.33570e-02,  7.32749e-01, -1.08837e-01,  ..., -5.22307e+00, -4.27569e+00, -3.46241e+00],\n",
       "             [ 1.34392e-01,  8.02712e-01,  4.78049e-01,  ..., -6.05456e+00, -3.96064e+00, -4.36903e+00],\n",
       "             [ 4.37943e-01,  5.46213e-01,  4.42128e-01,  ..., -7.01148e+00, -4.22282e+00, -5.05536e+00],\n",
       "             ...,\n",
       "             [ 1.37286e-01,  5.04759e-01,  3.40639e-01,  ..., -7.17813e+00, -3.88370e+00, -4.85044e+00],\n",
       "             [ 1.18137e-01,  8.34474e-01,  2.56156e-01,  ..., -7.75473e+00, -4.13991e+00, -4.68290e+00],\n",
       "             [ 6.39255e-02,  7.01134e-01, -2.00645e-01,  ..., -6.18465e+00, -4.70670e+00, -2.82985e+00]],\n",
       "  \n",
       "            [[ 9.41771e-02,  1.40582e+00, -5.74411e-02,  ..., -4.98257e+00, -4.95407e+00, -3.10712e+00],\n",
       "             [-3.66859e-01,  1.26341e+00,  4.30508e-01,  ..., -5.75668e+00, -4.46427e+00, -4.40536e+00],\n",
       "             [ 1.03081e-01,  1.31767e+00,  3.30656e-01,  ..., -6.84618e+00, -4.72957e+00, -5.24981e+00],\n",
       "             ...,\n",
       "             [-4.60769e-02,  1.63930e+00,  4.28108e-02,  ..., -7.93241e+00, -5.11174e+00, -4.93699e+00],\n",
       "             [ 6.14007e-01,  1.20655e+00,  3.07768e-01,  ..., -7.45741e+00, -4.95828e+00, -4.39455e+00],\n",
       "             [-2.72251e-01,  1.12794e+00,  4.99593e-02,  ..., -5.32914e+00, -4.33904e+00, -2.97891e+00]],\n",
       "  \n",
       "            [[ 1.38353e-01, -2.87053e-01, -1.14478e-01,  ..., -5.35581e+00, -5.79658e+00, -2.63652e+00],\n",
       "             [ 1.33081e-01, -2.12182e-01,  2.53560e-01,  ..., -5.86213e+00, -4.99327e+00, -3.02297e+00],\n",
       "             [-1.15661e-01, -4.54478e-03,  5.36640e-01,  ..., -5.62778e+00, -4.70847e+00, -2.65400e+00],\n",
       "             ...,\n",
       "             [-2.39891e-01, -1.66015e-01, -6.10571e-02,  ..., -8.29023e+00, -6.32362e+00, -4.13099e+00],\n",
       "             [-3.58946e-01, -4.76147e-02,  3.44187e-01,  ..., -6.35670e+00, -4.88263e+00, -2.44679e+00],\n",
       "             [-5.14127e-01,  2.32002e-02,  1.26552e-01,  ..., -5.21967e+00, -4.76472e+00, -2.10135e+00]]]]]),\n",
       "  tensor([[[[[-3.81936e-01,  5.06525e-02,  1.80529e-01,  ..., -7.60565e+00, -3.65650e+00, -6.66790e+00],\n",
       "             [ 7.43924e-02, -4.11132e-01,  7.44963e-01,  ..., -8.37370e+00, -3.29736e+00, -6.03957e+00],\n",
       "             [-4.04586e-01, -8.75358e-02,  4.64565e-01,  ..., -8.34608e+00, -4.37171e+00, -6.13352e+00],\n",
       "             ...,\n",
       "             [-2.13863e-01, -6.51577e-02,  4.89232e-01,  ..., -7.82991e+00, -4.24515e+00, -5.84589e+00],\n",
       "             [ 1.36291e-01, -2.90093e-01,  1.14062e-01,  ..., -8.07876e+00, -4.87834e+00, -6.14164e+00],\n",
       "             [-8.63814e-02, -5.37004e-01,  2.97235e-01,  ..., -6.21872e+00, -3.55495e+00, -5.65405e+00]],\n",
       "  \n",
       "            [[-5.56206e-01,  4.76484e-01, -8.79405e-02,  ..., -8.58296e+00, -4.59297e+00, -6.76305e+00],\n",
       "             [-4.78150e-02,  6.28532e-01,  8.94212e-01,  ..., -9.44206e+00, -3.52875e+00, -5.89151e+00],\n",
       "             [-4.09393e-01,  3.06928e-01,  6.68955e-01,  ..., -9.56785e+00, -3.88296e+00, -5.72063e+00],\n",
       "             ...,\n",
       "             [-1.31251e-01,  7.72461e-01,  2.66918e-01,  ..., -8.90712e+00, -5.54683e+00, -6.22047e+00],\n",
       "             [-2.92057e-01,  3.90195e-01,  1.73604e-02,  ..., -1.03047e+01, -6.84992e+00, -7.14211e+00],\n",
       "             [ 1.90702e-01,  3.76747e-02,  6.54567e-02,  ..., -7.14127e+00, -3.81097e+00, -5.58950e+00]],\n",
       "  \n",
       "            [[-4.11578e-01, -8.93881e-01,  4.08457e-02,  ..., -9.32524e+00, -3.40430e+00, -6.75359e+00],\n",
       "             [ 3.64861e-01,  6.04756e-01,  5.31444e-01,  ..., -1.09509e+01, -4.71524e+00, -6.09262e+00],\n",
       "             [ 3.85141e-01,  7.13554e-01,  5.89193e-01,  ..., -1.07877e+01, -4.82707e+00, -6.09378e+00],\n",
       "             ...,\n",
       "             [-3.49787e-01, -7.53765e-01,  2.04336e-01,  ..., -1.04523e+01, -5.68684e+00, -7.16908e+00],\n",
       "             [ 2.82434e-01, -5.28987e-01,  7.03085e-02,  ..., -1.02930e+01, -6.32420e+00, -7.06130e+00],\n",
       "             [ 4.03582e-01,  6.67328e-02,  1.55053e-02,  ..., -7.74277e+00, -4.14405e+00, -5.89111e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[-5.10157e-01,  2.74326e-01, -3.85201e-02,  ..., -8.50420e+00, -4.04925e+00, -6.02222e+00],\n",
       "             [ 5.99692e-01, -8.68693e-02,  5.05786e-01,  ..., -9.53417e+00, -4.80612e+00, -5.74712e+00],\n",
       "             [-3.40234e-01, -4.22963e-01,  6.74176e-01,  ..., -8.93342e+00, -4.23986e+00, -5.62406e+00],\n",
       "             ...,\n",
       "             [ 4.10472e-01,  5.91184e-01,  4.33687e-01,  ..., -1.01663e+01, -5.27344e+00, -6.21713e+00],\n",
       "             [ 3.26509e-01,  3.33455e-02,  3.67764e-01,  ..., -9.67266e+00, -4.80841e+00, -5.97412e+00],\n",
       "             [-2.93308e-02, -9.74558e-02,  2.40256e-01,  ..., -7.87440e+00, -3.79865e+00, -5.66293e+00]],\n",
       "  \n",
       "            [[-4.43839e-01,  4.12185e-01,  1.09350e-01,  ..., -7.33437e+00, -4.47751e+00, -5.84843e+00],\n",
       "             [ 1.20826e-02,  1.03946e+00,  6.49766e-01,  ..., -8.87840e+00, -4.49885e+00, -5.98713e+00],\n",
       "             [-8.85187e-01,  1.22463e+00,  7.33599e-01,  ..., -8.11075e+00, -3.45864e+00, -5.78059e+00],\n",
       "             ...,\n",
       "             [ 1.06987e-03, -4.73138e-01,  3.56558e-01,  ..., -1.01952e+01, -4.62880e+00, -6.14338e+00],\n",
       "             [ 2.89613e-01, -5.98608e-01,  5.99779e-01,  ..., -9.50962e+00, -4.71441e+00, -5.59107e+00],\n",
       "             [ 2.93227e-01,  8.12016e-01,  1.95569e-01,  ..., -7.49313e+00, -3.42719e+00, -5.88452e+00]],\n",
       "  \n",
       "            [[-3.38409e-01,  1.10201e-02,  2.04542e-01,  ..., -6.85034e+00, -4.55365e+00, -6.66957e+00],\n",
       "             [-1.51357e-01,  1.81212e-01,  8.58831e-01,  ..., -7.55065e+00, -3.47558e+00, -6.19967e+00],\n",
       "             [-7.08988e-01,  1.96001e-01,  6.95437e-01,  ..., -8.02311e+00, -3.31119e+00, -6.31010e+00],\n",
       "             ...,\n",
       "             [-1.32196e-01,  4.44246e-01,  4.37019e-01,  ..., -8.63171e+00, -3.99569e+00, -6.35304e+00],\n",
       "             [ 5.58681e-02,  2.85654e-01,  3.64476e-01,  ..., -8.60678e+00, -4.56003e+00, -6.55249e+00],\n",
       "             [-1.54245e-01,  3.09873e-01,  4.28349e-01,  ..., -5.84926e+00, -2.78396e+00, -5.92586e+00]]],\n",
       "  \n",
       "  \n",
       "           [[[-2.71518e-01, -4.93970e-02, -3.25622e-01,  ..., -5.76258e+00, -2.88301e+00, -3.07375e+00],\n",
       "             [ 1.95057e-01, -5.08742e-01,  3.33759e-01,  ..., -5.31921e+00, -2.78093e+00, -4.77073e+00],\n",
       "             [-4.16753e-01, -1.77829e-01, -3.07981e-02,  ..., -5.35357e+00, -3.06170e+00, -5.09822e+00],\n",
       "             ...,\n",
       "             [-1.98463e-01, -1.93950e-01,  1.55261e-01,  ..., -5.05508e+00, -2.59238e+00, -4.03970e+00],\n",
       "             [ 2.30014e-01, -4.21053e-01,  2.78638e-02,  ..., -5.39786e+00, -3.24254e+00, -3.33676e+00],\n",
       "             [-1.90733e-01, -6.46729e-01, -4.94780e-02,  ..., -5.55557e+00, -3.09308e+00, -2.49335e+00]],\n",
       "  \n",
       "            [[-5.31021e-01,  4.18989e-01, -5.95761e-01,  ..., -5.84695e+00, -3.46998e+00, -3.55365e+00],\n",
       "             [-1.21932e-01,  4.36947e-01,  2.57818e-01,  ..., -4.92064e+00, -2.61220e+00, -7.45608e+00],\n",
       "             [-5.56217e-01,  6.73174e-02,  7.83416e-02,  ..., -4.80785e+00, -2.93619e+00, -6.96018e+00],\n",
       "             ...,\n",
       "             [-1.33937e-01,  6.79613e-01, -1.12597e-01,  ..., -5.24488e+00, -3.73599e+00, -4.33365e+00],\n",
       "             [-2.69982e-01,  3.29080e-01, -3.95366e-01,  ..., -5.78886e+00, -4.18283e+00, -4.04521e+00],\n",
       "             [ 1.13583e-01, -7.97346e-02, -1.99056e-01,  ..., -5.16151e+00, -3.08830e+00, -3.85938e+00]],\n",
       "  \n",
       "            [[-3.97757e-01, -9.58053e-01, -4.76418e-01,  ..., -5.85483e+00, -2.26919e+00, -4.94615e+00],\n",
       "             [ 3.58382e-01,  2.57644e-01,  6.22667e-03,  ..., -5.15954e+00, -2.36748e+00, -7.33040e+00],\n",
       "             [ 3.79552e-01,  4.87149e-01,  1.16704e-03,  ..., -4.90737e+00, -2.53234e+00, -7.28402e+00],\n",
       "             ...,\n",
       "             [-3.48060e-01, -8.22532e-01, -2.34337e-01,  ..., -6.08657e+00, -3.39436e+00, -4.67537e+00],\n",
       "             [ 3.15850e-01, -5.87603e-01, -3.64095e-01,  ..., -5.93037e+00, -4.03037e+00, -3.67927e+00],\n",
       "             [ 3.58508e-01, -1.64641e-01, -3.84322e-01,  ..., -5.59067e+00, -3.41194e+00, -3.28343e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[-4.74885e-01,  1.54623e-01, -4.93195e-01,  ..., -5.27428e+00, -3.14731e+00, -5.37337e+00],\n",
       "             [ 5.69261e-01, -2.00653e-01,  3.92617e-02,  ..., -4.89628e+00, -3.30819e+00, -7.09338e+00],\n",
       "             [-4.08688e-01, -4.90526e-01,  8.18474e-02,  ..., -4.83187e+00, -3.14894e+00, -6.70966e+00],\n",
       "             ...,\n",
       "             [ 4.31337e-01,  5.46453e-01, -6.91020e-02,  ..., -4.97708e+00, -2.90016e+00, -7.12407e+00],\n",
       "             [ 3.85321e-01, -2.04711e-02, -6.39038e-02,  ..., -4.83523e+00, -3.06321e+00, -6.80050e+00],\n",
       "             [-1.09652e-01, -2.05845e-01, -1.41788e-01,  ..., -5.18540e+00, -3.10584e+00, -5.33911e+00]],\n",
       "  \n",
       "            [[-3.08580e-01,  2.38957e-01, -2.94498e-01,  ..., -5.26668e+00, -4.08857e+00, -3.64674e+00],\n",
       "             [-8.14872e-02,  9.61583e-01,  1.47962e-01,  ..., -4.86969e+00, -3.66924e+00, -5.96277e+00],\n",
       "             [-9.78621e-01,  1.07813e+00,  6.28102e-02,  ..., -4.60968e+00, -2.71925e+00, -5.85066e+00],\n",
       "             ...,\n",
       "             [ 2.67737e-02, -5.21323e-01, -7.63511e-02,  ..., -5.13988e+00, -3.05182e+00, -6.79442e+00],\n",
       "             [ 3.47295e-01, -6.93958e-01,  1.59644e-01,  ..., -5.12214e+00, -3.35206e+00, -6.27104e+00],\n",
       "             [ 2.24976e-01,  6.86121e-01, -2.62975e-01,  ..., -5.24495e+00, -2.84691e+00, -5.15599e+00]],\n",
       "  \n",
       "            [[-9.76256e-02,  2.32869e-02, -1.50090e-01,  ..., -5.80129e+00, -4.05384e+00, -2.23485e+00],\n",
       "             [-8.18976e-02,  1.09436e-01,  4.75617e-01,  ..., -5.23208e+00, -3.25446e+00, -4.30723e+00],\n",
       "             [-8.48546e-01,  1.14902e-01,  2.09862e-01,  ..., -5.24746e+00, -2.80424e+00, -4.46853e+00],\n",
       "             ...,\n",
       "             [-1.45730e-01,  3.78047e-01,  3.93969e-02,  ..., -5.33908e+00, -2.65316e+00, -4.82650e+00],\n",
       "             [ 1.65160e-01,  2.10417e-01,  3.65645e-03,  ..., -5.57361e+00, -2.94596e+00, -3.77712e+00],\n",
       "             [-2.26357e-01,  2.39879e-01, -1.06218e-02,  ..., -5.53466e+00, -2.95309e+00, -3.16548e+00]]],\n",
       "  \n",
       "  \n",
       "           [[[-2.57038e-01,  1.05213e-01, -1.51579e-01,  ..., -5.80642e+00, -2.72167e+00, -4.21037e+00],\n",
       "             [ 2.00948e-01, -3.55343e-01,  6.59535e-01,  ..., -6.34746e+00, -2.33368e+00, -4.87361e+00],\n",
       "             [-3.91001e-01, -6.84696e-02,  2.00441e-01,  ..., -6.49714e+00, -3.26354e+00, -5.34834e+00],\n",
       "             ...,\n",
       "             [-1.83472e-01, -4.24981e-05,  4.65496e-01,  ..., -6.09241e+00, -3.00880e+00, -5.76451e+00],\n",
       "             [ 2.38286e-01, -2.44246e-01,  3.10412e-01,  ..., -5.67640e+00, -3.56957e+00, -4.05399e+00],\n",
       "             [-1.87006e-01, -4.96385e-01,  2.23914e-01,  ..., -4.46117e+00, -2.43321e+00, -2.99880e+00]],\n",
       "  \n",
       "            [[-5.16957e-01,  4.55324e-01, -4.55836e-01,  ..., -6.41801e+00, -3.34265e+00, -4.35179e+00],\n",
       "             [-8.99182e-02,  5.05415e-01,  5.81378e-01,  ..., -7.72496e+00, -2.34600e+00, -7.08567e+00],\n",
       "             [-5.23359e-01,  9.43163e-02,  3.80807e-01,  ..., -7.68864e+00, -2.50833e+00, -8.42573e+00],\n",
       "             ...,\n",
       "             [-1.06574e-01,  6.80535e-01,  1.44545e-01,  ..., -6.80765e+00, -3.63066e+00, -6.84322e+00],\n",
       "             [-2.50987e-01,  3.66742e-01, -2.27197e-01,  ..., -7.64755e+00, -5.37004e+00, -4.92280e+00],\n",
       "             [ 1.10430e-01,  6.22558e-02,  9.02129e-02,  ..., -5.50834e+00, -2.77486e+00, -3.99504e+00]],\n",
       "  \n",
       "            [[-3.84737e-01, -9.04399e-01, -3.34679e-01,  ..., -7.69164e+00, -2.20161e+00, -6.17236e+00],\n",
       "             [ 3.83222e-01,  3.63955e-01,  2.28084e-01,  ..., -8.74595e+00, -2.79267e+00, -8.67870e+00],\n",
       "             [ 3.97868e-01,  5.68577e-01,  2.17830e-01,  ..., -8.74241e+00, -3.02096e+00, -9.76587e+00],\n",
       "             ...,\n",
       "             [-3.23253e-01, -8.15711e-01, -5.41075e-02,  ..., -8.47847e+00, -4.02703e+00, -6.56153e+00],\n",
       "             [ 3.24666e-01, -5.65423e-01, -1.85392e-01,  ..., -7.80354e+00, -4.85351e+00, -5.08007e+00],\n",
       "             [ 3.68889e-01, -1.70078e-02, -1.60677e-01,  ..., -5.81379e+00, -2.85980e+00, -3.87037e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[-4.57598e-01,  1.91393e-01, -3.61808e-01,  ..., -6.52282e+00, -2.81391e+00, -5.18060e+00],\n",
       "             [ 5.84594e-01, -1.61110e-01,  3.15219e-01,  ..., -7.42761e+00, -3.60210e+00, -7.16518e+00],\n",
       "             [-3.79123e-01, -4.57034e-01,  3.42245e-01,  ..., -7.42775e+00, -3.38336e+00, -7.97965e+00],\n",
       "             ...,\n",
       "             [ 4.55275e-01,  5.53312e-01,  1.69575e-01,  ..., -8.41730e+00, -3.79901e+00, -8.63463e+00],\n",
       "             [ 4.04907e-01, -8.82536e-04,  1.86230e-01,  ..., -8.30607e+00, -3.78338e+00, -7.27480e+00],\n",
       "             [-9.74789e-02, -1.75475e-01,  1.37701e-01,  ..., -6.57461e+00, -2.97305e+00, -4.72359e+00]],\n",
       "  \n",
       "            [[-2.95135e-01,  3.20300e-01, -1.52340e-01,  ..., -5.50352e+00, -3.67702e+00, -3.30600e+00],\n",
       "             [-6.03274e-02,  1.03379e+00,  4.15698e-01,  ..., -6.81076e+00, -3.77806e+00, -5.87888e+00],\n",
       "             [-9.49324e-01,  1.17302e+00,  2.88216e-01,  ..., -6.83586e+00, -2.88795e+00, -7.28605e+00],\n",
       "             ...,\n",
       "             [ 4.77894e-02, -4.56994e-01,  1.44445e-01,  ..., -8.59881e+00, -3.65468e+00, -6.85802e+00],\n",
       "             [ 3.65304e-01, -6.29815e-01,  4.37880e-01,  ..., -7.86041e+00, -3.63547e+00, -6.52155e+00],\n",
       "             [ 2.24700e-01,  8.01689e-01, -5.25795e-02,  ..., -6.53102e+00, -2.74271e+00, -4.61725e+00]],\n",
       "  \n",
       "            [[-9.60310e-02,  2.25310e-02,  6.14090e-03,  ..., -4.80758e+00, -3.81029e+00, -1.48403e+00],\n",
       "             [-6.81846e-02,  9.70174e-02,  8.54552e-01,  ..., -5.70557e+00, -2.59952e+00, -4.09308e+00],\n",
       "             [-8.26091e-01,  1.04970e-01,  4.84425e-01,  ..., -6.20606e+00, -2.53026e+00, -4.28361e+00],\n",
       "             ...,\n",
       "             [-1.38679e-01,  3.68963e-01,  3.52652e-01,  ..., -6.74856e+00, -2.64076e+00, -5.29916e+00],\n",
       "             [ 1.78774e-01,  1.80212e-01,  3.06511e-01,  ..., -6.76685e+00, -3.20822e+00, -3.53471e+00],\n",
       "             [-2.14162e-01,  2.03329e-01,  2.66928e-01,  ..., -4.92588e+00, -1.91244e+00, -2.67526e+00]]]]]),\n",
       "  tensor([[[[[-4.23741e-01, -1.09785e-01, -6.12193e-01,  ..., -5.23550e+00, -3.42107e+00, -2.73256e+00],\n",
       "             [-9.27796e-01, -3.20770e-02, -3.39327e-01,  ..., -5.43666e+00, -5.44380e+00, -2.66688e+00],\n",
       "             [-5.16473e-02, -4.01275e-01, -2.18551e-01,  ..., -5.72541e+00, -6.21033e+00, -3.02914e+00],\n",
       "             ...,\n",
       "             [-8.34140e-01,  1.19980e-01, -3.88781e-01,  ..., -5.48191e+00, -6.24741e+00, -2.52155e+00],\n",
       "             [ 8.03812e-02,  2.14119e-01, -2.06766e-01,  ..., -5.48771e+00, -5.62617e+00, -3.23844e+00],\n",
       "             [-7.75238e-02,  6.08373e-01, -3.78971e-01,  ..., -5.18783e+00, -3.46805e+00, -2.84012e+00]],\n",
       "  \n",
       "            [[-2.20763e-01, -1.89761e-01, -7.31181e-01,  ..., -5.63365e+00, -5.22865e+00, -2.84911e+00],\n",
       "             [ 9.39810e-02,  1.95350e-01, -5.87716e-01,  ..., -6.30725e+00, -7.49675e+00, -3.83752e+00],\n",
       "             [ 1.07893e-01, -3.49913e-01, -3.49756e-01,  ..., -6.21514e+00, -8.22048e+00, -3.93486e+00],\n",
       "             ...,\n",
       "             [-6.85597e-01, -2.43843e-01, -6.81652e-01,  ..., -6.27139e+00, -7.76589e+00, -3.58342e+00],\n",
       "             [ 1.22149e-01,  4.14371e-02, -5.81175e-01,  ..., -6.06898e+00, -6.92914e+00, -3.48520e+00],\n",
       "             [-2.14903e-01, -4.35685e-01, -6.75228e-01,  ..., -5.68782e+00, -5.64809e+00, -2.71591e+00]],\n",
       "  \n",
       "            [[-6.41645e-01, -1.44947e-02, -8.21920e-01,  ..., -5.71494e+00, -5.66172e+00, -1.92199e+00],\n",
       "             [-2.27855e-01, -5.05580e-02, -4.92199e-01,  ..., -6.05239e+00, -6.44727e+00, -2.70514e+00],\n",
       "             [-4.08678e-01, -2.03084e-01, -3.75160e-01,  ..., -6.58326e+00, -8.33224e+00, -2.85437e+00],\n",
       "             ...,\n",
       "             [-5.11617e-01, -7.80379e-01, -5.61225e-01,  ..., -6.00210e+00, -6.96742e+00, -3.35854e+00],\n",
       "             [ 2.92395e-01,  2.66555e-01, -5.11676e-01,  ..., -6.17440e+00, -6.35762e+00, -3.32443e+00],\n",
       "             [-3.54713e-01, -8.78915e-02, -6.14859e-01,  ..., -5.36413e+00, -4.23964e+00, -2.46134e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[ 1.00648e-01,  1.88343e-01, -6.95496e-01,  ..., -5.62886e+00, -4.70916e+00, -2.75962e+00],\n",
       "             [ 2.16074e-01, -5.77496e-01, -6.00421e-01,  ..., -6.38832e+00, -7.46881e+00, -2.59222e+00],\n",
       "             [-4.56567e-01, -8.73271e-01, -4.17258e-01,  ..., -6.26300e+00, -7.74584e+00, -2.66290e+00],\n",
       "             ...,\n",
       "             [ 2.69317e-01, -8.78602e-01, -6.18295e-01,  ..., -6.71448e+00, -7.49538e+00, -3.42825e+00],\n",
       "             [-1.21650e+00, -4.03241e-01, -4.05979e-01,  ..., -6.20497e+00, -7.21574e+00, -4.13745e+00],\n",
       "             [ 7.51036e-01,  1.28753e-02, -5.29622e-01,  ..., -5.75380e+00, -4.70307e+00, -3.55622e+00]],\n",
       "  \n",
       "            [[-6.98189e-01,  2.51890e-01, -7.56177e-01,  ..., -5.64539e+00, -4.31522e+00, -2.88006e+00],\n",
       "             [-1.58705e-01, -1.23460e-01, -8.38019e-02,  ..., -5.44983e+00, -6.34869e+00, -3.56691e+00],\n",
       "             [ 1.05670e+00, -1.75001e-01, -5.16644e-01,  ..., -6.28450e+00, -7.17898e+00, -3.92195e+00],\n",
       "             ...,\n",
       "             [ 6.17517e-01, -1.08161e+00, -3.77318e-01,  ..., -5.94121e+00, -6.21792e+00, -4.05245e+00],\n",
       "             [ 5.28627e-01, -3.80428e-01, -5.77852e-02,  ..., -5.61097e+00, -6.40478e+00, -4.58823e+00],\n",
       "             [ 4.20282e-01, -3.52398e-01, -1.88810e-01,  ..., -5.32835e+00, -4.25402e+00, -3.73582e+00]],\n",
       "  \n",
       "            [[-8.76454e-02,  2.72754e-01, -1.26666e-01,  ..., -4.82950e+00, -3.92103e+00, -1.94649e+00],\n",
       "             [-1.33078e+00, -1.34499e-01,  6.62495e-02,  ..., -4.58915e+00, -4.67583e+00, -2.30133e+00],\n",
       "             [ 3.76911e-01, -5.09528e-01, -2.55008e-01,  ..., -4.93801e+00, -5.75169e+00, -2.42367e+00],\n",
       "             ...,\n",
       "             [ 3.62009e-01,  5.41807e-01, -2.05906e-01,  ..., -5.08291e+00, -4.79870e+00, -2.76401e+00],\n",
       "             [-2.16091e-01,  6.98583e-01,  1.29072e-01,  ..., -4.59282e+00, -3.98749e+00, -2.96616e+00],\n",
       "             [-3.55756e-01,  3.33767e-01, -6.31715e-03,  ..., -4.87023e+00, -4.30267e+00, -2.98805e+00]]],\n",
       "  \n",
       "  \n",
       "           [[[-4.38401e-01, -8.86997e-02, -5.38957e-01,  ..., -4.66787e+00, -4.87651e+00, -6.60928e+00],\n",
       "             [-9.23241e-01,  3.89066e-03, -3.29764e-01,  ..., -6.81854e+00, -6.55808e+00, -7.32352e+00],\n",
       "             [-4.55696e-02, -3.61993e-01, -1.86666e-01,  ..., -7.68591e+00, -7.28706e+00, -7.36245e+00],\n",
       "             ...,\n",
       "             [-8.29473e-01,  1.58963e-01, -4.41869e-01,  ..., -8.19667e+00, -7.36414e+00, -8.14309e+00],\n",
       "             [ 9.01312e-02,  2.55061e-01, -1.61138e-01,  ..., -6.71474e+00, -6.42576e+00, -7.67107e+00],\n",
       "             [-6.46439e-02,  6.41552e-01, -3.20168e-01,  ..., -4.25908e+00, -4.35535e+00, -6.72813e+00]],\n",
       "  \n",
       "            [[-2.26324e-01, -2.08872e-01, -7.15958e-01,  ..., -7.14276e+00, -6.73021e+00, -7.58266e+00],\n",
       "             [ 9.54524e-02,  1.96799e-01, -6.34208e-01,  ..., -1.08872e+01, -8.90348e+00, -8.67582e+00],\n",
       "             [ 1.19054e-01, -3.38277e-01, -3.45215e-01,  ..., -1.08335e+01, -8.88584e+00, -8.20774e+00],\n",
       "             ...,\n",
       "             [-6.84266e-01, -2.50752e-01, -7.56549e-01,  ..., -1.17068e+01, -9.12899e+00, -9.41522e+00],\n",
       "             [ 1.37314e-01,  7.37617e-02, -6.32586e-01,  ..., -9.39127e+00, -7.57747e+00, -8.30738e+00],\n",
       "             [-2.06744e-01, -4.21291e-01, -6.80420e-01,  ..., -6.57454e+00, -6.43860e+00, -7.62320e+00]],\n",
       "  \n",
       "            [[-6.52919e-01, -2.84067e-02, -8.35707e-01,  ..., -7.11593e+00, -6.44130e+00, -7.10682e+00],\n",
       "             [-2.29867e-01, -4.26270e-02, -5.56687e-01,  ..., -9.87188e+00, -7.95270e+00, -7.44244e+00],\n",
       "             [-3.98411e-01, -1.87081e-01, -4.63689e-01,  ..., -1.19899e+01, -9.48485e+00, -7.59702e+00],\n",
       "             ...,\n",
       "             [-5.10162e-01, -7.51485e-01, -7.17956e-01,  ..., -1.08749e+01, -8.02742e+00, -8.43666e+00],\n",
       "             [ 3.06110e-01,  2.97675e-01, -6.10627e-01,  ..., -1.05518e+01, -7.84038e+00, -8.41835e+00],\n",
       "             [-3.44105e-01, -6.09516e-02, -6.74170e-01,  ..., -7.05920e+00, -6.38160e+00, -7.72632e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[ 1.02151e-01,  2.07485e-01, -6.84929e-01,  ..., -6.73393e+00, -5.78494e+00, -6.64581e+00],\n",
       "             [ 2.14636e-01, -5.53661e-01, -6.15530e-01,  ..., -9.88179e+00, -8.89626e+00, -7.66565e+00],\n",
       "             [-4.56001e-01, -8.31095e-01, -4.64468e-01,  ..., -9.96118e+00, -8.15689e+00, -7.54379e+00],\n",
       "             ...,\n",
       "             [ 2.70157e-01, -9.11057e-01, -6.56387e-01,  ..., -1.19932e+01, -8.76442e+00, -8.37549e+00],\n",
       "             [-1.20453e+00, -3.91852e-01, -4.01560e-01,  ..., -1.16468e+01, -9.30167e+00, -8.91248e+00],\n",
       "             [ 7.70158e-01,  5.38788e-02, -4.79536e-01,  ..., -6.76434e+00, -6.33692e+00, -7.83447e+00]],\n",
       "  \n",
       "            [[-7.13748e-01,  2.63490e-01, -6.91554e-01,  ..., -6.28988e+00, -5.59456e+00, -6.93586e+00],\n",
       "             [-1.53979e-01, -1.07688e-01,  3.49565e-03,  ..., -8.63666e+00, -7.86480e+00, -7.14227e+00],\n",
       "             [ 1.05808e+00, -1.32260e-01, -5.37619e-01,  ..., -1.14877e+01, -9.02706e+00, -8.22731e+00],\n",
       "             ...,\n",
       "             [ 6.19725e-01, -1.06366e+00, -4.47197e-01,  ..., -1.10565e+01, -7.91930e+00, -7.68752e+00],\n",
       "             [ 5.39577e-01, -3.52065e-01, -7.38316e-02,  ..., -9.94268e+00, -8.26062e+00, -7.97732e+00],\n",
       "             [ 4.46019e-01, -3.46489e-01, -1.22211e-01,  ..., -6.16799e+00, -5.71836e+00, -7.01869e+00]],\n",
       "  \n",
       "            [[-9.43871e-02,  2.61428e-01, -2.41374e-02,  ..., -3.69671e+00, -4.68597e+00, -5.08581e+00],\n",
       "             [-1.31818e+00, -1.48053e-01,  1.89575e-01,  ..., -5.38269e+00, -5.91208e+00, -6.23892e+00],\n",
       "             [ 3.81308e-01, -5.37009e-01, -2.34194e-01,  ..., -6.80641e+00, -6.60389e+00, -6.29252e+00],\n",
       "             ...,\n",
       "             [ 3.69655e-01,  5.44495e-01, -2.15066e-01,  ..., -5.52310e+00, -5.11154e+00, -6.18338e+00],\n",
       "             [-2.03741e-01,  6.80166e-01,  2.25042e-01,  ..., -4.48221e+00, -4.89170e+00, -6.23735e+00],\n",
       "             [-3.38024e-01,  3.09545e-01,  1.25501e-01,  ..., -3.97349e+00, -4.71128e+00, -5.88328e+00]]],\n",
       "  \n",
       "  \n",
       "           [[[-2.88239e-01, -1.54364e-01, -7.68859e-01,  ..., -4.43759e+00, -2.96045e+00, -4.56551e+00],\n",
       "             [-8.72081e-01, -9.30171e-02, -4.46194e-01,  ..., -5.56871e+00, -4.19253e+00, -5.56650e+00],\n",
       "             [ 2.96191e-03, -4.09222e-01, -3.13835e-01,  ..., -6.05468e+00, -4.59788e+00, -5.92128e+00],\n",
       "             ...,\n",
       "             [-7.69190e-01,  5.52534e-02, -4.04996e-01,  ..., -6.01350e+00, -3.92328e+00, -6.80541e+00],\n",
       "             [ 1.07666e-01,  1.67246e-01, -2.51145e-01,  ..., -5.85730e+00, -3.71870e+00, -6.34788e+00],\n",
       "             [-1.05810e-01,  5.57585e-01, -5.04600e-01,  ..., -4.56678e+00, -2.59908e+00, -4.86332e+00]],\n",
       "  \n",
       "            [[-7.78846e-02, -3.59942e-01, -7.41451e-01,  ..., -5.62521e+00, -3.85149e+00, -5.59122e+00],\n",
       "             [ 1.62412e-01,  1.32359e-01, -5.66128e-01,  ..., -7.52390e+00, -5.12169e+00, -7.30104e+00],\n",
       "             [ 1.39964e-01, -3.94074e-01, -3.63657e-01,  ..., -8.01671e+00, -5.37047e+00, -7.32101e+00],\n",
       "             ...,\n",
       "             [-6.25746e-01, -2.97719e-01, -6.91426e-01,  ..., -7.25913e+00, -5.12949e+00, -8.05389e+00],\n",
       "             [ 1.13162e-01, -4.76129e-02, -4.80904e-01,  ..., -7.07521e+00, -4.27890e+00, -6.84481e+00],\n",
       "             [-2.59662e-01, -5.05459e-01, -7.36142e-01,  ..., -5.31163e+00, -4.04272e+00, -5.50741e+00]],\n",
       "  \n",
       "            [[-4.89395e-01, -1.50591e-01, -8.47921e-01,  ..., -5.51354e+00, -4.28162e+00, -5.02639e+00],\n",
       "             [-1.56334e-01, -2.28155e-01, -4.72111e-01,  ..., -7.53833e+00, -4.21360e+00, -6.52043e+00],\n",
       "             [-3.69429e-01, -3.34856e-01, -4.10586e-01,  ..., -9.14489e+00, -5.56076e+00, -7.19587e+00],\n",
       "             ...,\n",
       "             [-4.58538e-01, -8.64056e-01, -5.26119e-01,  ..., -7.95570e+00, -4.49437e+00, -7.83085e+00],\n",
       "             [ 2.93070e-01,  2.28258e-01, -4.33100e-01,  ..., -7.40295e+00, -4.04672e+00, -7.06531e+00],\n",
       "             [-3.99732e-01, -2.22038e-01, -6.25056e-01,  ..., -5.65357e+00, -2.75868e+00, -6.35821e+00]],\n",
       "  \n",
       "            ...,\n",
       "  \n",
       "            [[ 2.49529e-01,  9.18684e-02, -6.79018e-01,  ..., -5.95736e+00, -3.53430e+00, -4.53786e+00],\n",
       "             [ 2.80306e-01, -6.29131e-01, -6.80177e-01,  ..., -6.97977e+00, -5.72590e+00, -5.61401e+00],\n",
       "             [-3.98830e-01, -9.11427e-01, -4.52365e-01,  ..., -7.46676e+00, -5.31102e+00, -5.96437e+00],\n",
       "             ...,\n",
       "             [ 2.86991e-01, -1.00704e+00, -7.25274e-01,  ..., -8.04211e+00, -5.29579e+00, -7.41400e+00],\n",
       "             [-1.22588e+00, -4.31431e-01, -5.15553e-01,  ..., -7.63498e+00, -5.30886e+00, -7.33969e+00],\n",
       "             [ 6.64977e-01,  2.60141e-02, -6.23367e-01,  ..., -5.17256e+00, -3.49858e+00, -5.48359e+00]],\n",
       "  \n",
       "            [[-5.54858e-01,  2.02592e-01, -8.64175e-01,  ..., -5.02868e+00, -3.60762e+00, -4.60854e+00],\n",
       "             [-8.54575e-02, -2.11048e-01, -2.08280e-01,  ..., -6.57134e+00, -4.29843e+00, -5.73063e+00],\n",
       "             [ 1.11110e+00, -1.82949e-01, -4.95213e-01,  ..., -7.30207e+00, -4.96496e+00, -6.87053e+00],\n",
       "             ...,\n",
       "             [ 6.40117e-01, -1.19730e+00, -4.28173e-01,  ..., -7.52255e+00, -4.03154e+00, -7.04789e+00],\n",
       "             [ 5.23519e-01, -4.36266e-01, -1.32567e-01,  ..., -7.18846e+00, -4.32633e+00, -7.05301e+00],\n",
       "             [ 2.98548e-01, -3.63500e-01, -3.67465e-01,  ..., -5.28525e+00, -3.38763e+00, -5.18294e+00]],\n",
       "  \n",
       "            [[-1.30855e-02,  2.01924e-01, -4.92099e-01,  ..., -4.21433e+00, -3.96314e+00, -2.72808e+00],\n",
       "             [-1.28889e+00, -1.94373e-01, -2.28098e-01,  ..., -4.87629e+00, -4.27156e+00, -4.12740e+00],\n",
       "             [ 4.19096e-01, -5.79697e-01, -4.22980e-01,  ..., -5.58043e+00, -4.80853e+00, -4.87547e+00],\n",
       "             ...,\n",
       "             [ 3.86143e-01,  4.64154e-01, -3.93388e-01,  ..., -4.84867e+00, -4.20617e+00, -4.72161e+00],\n",
       "             [-2.10460e-01,  5.54236e-01, -1.09425e-01,  ..., -4.57047e+00, -3.47207e+00, -4.66363e+00],\n",
       "             [-4.30626e-01,  1.95835e-01, -3.05449e-01,  ..., -4.58198e+00, -3.75030e+00, -3.76724e+00]]]]])])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_model(torch.rand((1,3,384,672)).cpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "edd92dc4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: [Torch-TensorRT] - Cannot infer input type from calcuations in graph for input x.1. Assuming it is Float32. If not, specify input type explicity\n",
      "WARNING: [Torch-TensorRT] - Dilation not used in Max pooling converter\n",
      "WARNING: [Torch-TensorRT] - Dilation not used in Max pooling converter\n",
      "WARNING: [Torch-TensorRT] - Dilation not used in Max pooling converter\n",
      "WARNING: [Torch-TensorRT] - Dilation not used in Max pooling converter\n",
      "WARNING: [Torch-TensorRT] - Dilation not used in Max pooling converter\n",
      "WARNING: [Torch-TensorRT] - Dilation not used in Max pooling converter\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n",
      "WARNING: [Torch-TensorRT] - There may be undefined behavior using dynamic shape and aten::size\n"
     ]
    }
   ],
   "source": [
    "qat_model = torch.jit.load(\"trained_vgg16_qat.jit.pt\").eval()\n",
    "compile_spec = {\"inputs\": [torch_tensorrt.Input([1, 3, 384, 672])],\n",
    "                \"enabled_precisions\": torch.int8,\n",
    "                \"truncate_long_and_double\": True,\n",
    "                }\n",
    "trt_mod = torch_tensorrt.compile(qat_model, **compile_spec)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9c7d0a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_image = torch.rand((1,3,384,672)).cuda()\n",
    "o1 = qat_model(input_image)\n",
    "# o2 = jit_model(input_image)\n",
    "o3 = trt_mod(input_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1769b63b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.testing.assert_almost_equal(o1[0].cpu().numpy(),o3[0].cpu().numpy())\n",
    "# np.testing.assert_almost_equal(o1[1].cpu().numpy(),o2[1].cpu().numpy())\n",
    "# np.testing.assert_almost_equal(o1[2].cpu().numpy(),o2[2].cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "fa2a6982",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.jit.save(trt_mod, w.replace('.pt', '_qat.trt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "40867691",
   "metadata": {},
   "outputs": [],
   "source": [
    "trt_mod = torch.jit.load(w.replace('.pt', '_qat.trt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "9155d454",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[[[-2.22109e-01,  1.50951e-01,  5.21803e-01,  ..., -5.93185e+00, -5.16824e+00, -2.88184e+00],\n",
       "            [-5.07687e-01,  7.32737e-02,  6.19502e-01,  ..., -6.55571e+00, -5.13715e+00, -3.13503e+00],\n",
       "            [-2.99645e-01,  1.55214e-01,  5.13571e-01,  ..., -7.34297e+00, -5.36995e+00, -4.05769e+00],\n",
       "            ...,\n",
       "            [-3.03014e-01,  2.90084e-02,  6.87270e-01,  ..., -7.07415e+00, -5.00878e+00, -4.26948e+00],\n",
       "            [-3.77576e-01,  1.02184e-01,  7.82256e-01,  ..., -7.54790e+00, -6.13167e+00, -5.22417e+00],\n",
       "            [-3.26356e-01, -1.61818e-01,  8.72771e-01,  ..., -5.59588e+00, -5.10294e+00, -3.18662e+00]],\n",
       " \n",
       "           [[-1.73041e-01, -9.29553e-01,  2.76176e-01,  ..., -6.00224e+00, -5.54974e+00, -3.42740e+00],\n",
       "            [-6.80190e-01, -1.48092e-01,  5.99334e-01,  ..., -7.14898e+00, -5.58389e+00, -4.51367e+00],\n",
       "            [-3.38186e-01, -4.56995e-01,  6.28369e-01,  ..., -7.10099e+00, -5.24368e+00, -4.00169e+00],\n",
       "            ...,\n",
       "            [-3.99954e-01,  1.99067e-01,  6.43019e-01,  ..., -8.11969e+00, -6.11672e+00, -5.97386e+00],\n",
       "            [-5.07719e-01, -8.06403e-01,  7.28333e-01,  ..., -7.57747e+00, -6.32625e+00, -4.97483e+00],\n",
       "            [ 1.43894e-01, -2.12454e-01,  7.22360e-01,  ..., -6.31040e+00, -5.58844e+00, -4.48859e+00]],\n",
       " \n",
       "           [[-3.85886e-01, -2.94949e-01,  2.68029e-01,  ..., -6.26605e+00, -5.16047e+00, -3.46388e+00],\n",
       "            [-1.00265e+00, -5.84385e-01,  7.13084e-01,  ..., -6.67954e+00, -4.98440e+00, -3.81943e+00],\n",
       "            [-3.43794e-01, -1.66636e-01,  7.89039e-01,  ..., -6.40571e+00, -4.68996e+00, -3.65262e+00],\n",
       "            ...,\n",
       "            [-1.05620e-01, -6.34467e-01,  6.15187e-01,  ..., -8.59107e+00, -5.90719e+00, -5.91000e+00],\n",
       "            [-4.98241e-02, -4.25369e-02,  5.87587e-01,  ..., -9.16211e+00, -6.81054e+00, -7.02496e+00],\n",
       "            [ 5.34456e-02, -4.31627e-01,  5.88211e-01,  ..., -7.08407e+00, -5.99086e+00, -4.85667e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-7.05031e-01, -3.29473e-01,  5.77369e-01,  ..., -4.87984e+00, -3.72479e+00, -2.96655e+00],\n",
       "            [-2.64769e-01, -1.48293e-01,  1.12112e+00,  ..., -5.50537e+00, -3.92976e+00, -3.33273e+00],\n",
       "            [-2.92016e-01, -3.36744e-01,  7.82306e-01,  ..., -6.24755e+00, -4.29561e+00, -3.58738e+00],\n",
       "            ...,\n",
       "            [-1.23555e+00,  5.02766e-01,  5.78668e-01,  ..., -7.05280e+00, -5.20358e+00, -4.61582e+00],\n",
       "            [-3.56407e-01,  2.63073e-01,  5.65867e-01,  ..., -7.24308e+00, -5.30983e+00, -4.13331e+00],\n",
       "            [-4.28834e-02,  2.82916e-02,  5.83620e-01,  ..., -5.55000e+00, -4.84236e+00, -2.28560e+00]],\n",
       " \n",
       "           [[-4.58783e-01,  6.78977e-01,  7.37717e-01,  ..., -4.99157e+00, -4.50563e+00, -2.83804e+00],\n",
       "            [-8.39545e-01,  9.65925e-01,  1.18220e+00,  ..., -5.21137e+00, -4.37957e+00, -3.00259e+00],\n",
       "            [-4.18152e-01,  9.64495e-01,  7.56288e-01,  ..., -6.38040e+00, -5.46177e+00, -3.78457e+00],\n",
       "            ...,\n",
       "            [-6.05890e-01,  7.05672e-01,  4.74363e-01,  ..., -7.48802e+00, -6.18194e+00, -5.32062e+00],\n",
       "            [-1.05681e+00,  9.86160e-01,  7.36811e-01,  ..., -6.85776e+00, -5.48894e+00, -3.36950e+00],\n",
       "            [ 3.07897e-01,  3.76802e-01,  3.42808e-01,  ..., -6.30861e+00, -5.79318e+00, -3.65886e+00]],\n",
       " \n",
       "           [[-5.59072e-01, -3.63140e-01,  7.74364e-01,  ..., -4.95387e+00, -5.49544e+00, -1.43400e+00],\n",
       "            [-7.48380e-01,  6.77812e-02,  1.28397e+00,  ..., -5.14504e+00, -5.01747e+00, -2.03101e+00],\n",
       "            [-6.02395e-01, -3.86344e-02,  8.85032e-01,  ..., -6.35225e+00, -5.77438e+00, -2.60894e+00],\n",
       "            ...,\n",
       "            [-5.82800e-01, -4.27389e-01,  5.73952e-01,  ..., -6.92064e+00, -6.19405e+00, -3.84861e+00],\n",
       "            [-7.87591e-01, -1.05059e-01,  5.59044e-01,  ..., -6.86640e+00, -5.82971e+00, -2.74797e+00],\n",
       "            [-9.28560e-02, -6.56014e-01,  4.14481e-01,  ..., -6.16259e+00, -5.47015e+00, -3.02869e+00]]],\n",
       " \n",
       " \n",
       "          [[[ 4.76212e-02,  3.50313e-01,  4.36005e-01,  ..., -5.69046e+00, -5.09209e+00, -2.25269e+00],\n",
       "            [ 5.37974e-01,  1.79207e-01,  5.88338e-01,  ..., -6.75047e+00, -5.48281e+00, -2.33910e+00],\n",
       "            [ 8.92268e-01,  3.74533e-01,  4.00360e-01,  ..., -8.00586e+00, -5.45500e+00, -3.14441e+00],\n",
       "            ...,\n",
       "            [ 8.89103e-01,  2.90091e-01,  6.88351e-01,  ..., -7.56329e+00, -5.08778e+00, -3.45631e+00],\n",
       "            [ 4.28829e-01,  3.24041e-01,  5.53420e-01,  ..., -8.32184e+00, -6.36182e+00, -4.31729e+00],\n",
       "            [ 1.23266e-01,  5.33434e-02,  5.34073e-01,  ..., -5.40597e+00, -5.12029e+00, -2.53438e+00]],\n",
       " \n",
       "           [[-1.61577e-01, -4.04140e-01,  2.36338e-01,  ..., -6.30903e+00, -5.38318e+00, -3.11288e+00],\n",
       "            [ 3.72863e-01, -2.56091e-01,  5.43694e-01,  ..., -7.95942e+00, -6.02582e+00, -4.31803e+00],\n",
       "            [ 8.48875e-01, -3.77904e-01,  5.93167e-01,  ..., -7.85288e+00, -5.46195e+00, -3.49333e+00],\n",
       "            ...,\n",
       "            [ 8.91482e-01,  3.47098e-01,  5.24938e-01,  ..., -9.70966e+00, -6.43897e+00, -5.90663e+00],\n",
       "            [ 5.80841e-01, -3.78456e-01,  7.27699e-01,  ..., -8.69926e+00, -6.60702e+00, -4.71211e+00],\n",
       "            [ 2.19486e-01,  2.26130e-01,  4.54097e-01,  ..., -6.23874e+00, -5.45560e+00, -4.06999e+00]],\n",
       " \n",
       "           [[-2.18630e-01,  1.16172e-01,  2.50826e-01,  ..., -6.46123e+00, -4.98940e+00, -2.92293e+00],\n",
       "            [ 9.16880e-02, -4.82175e-01,  8.98644e-01,  ..., -7.06344e+00, -5.27182e+00, -3.69286e+00],\n",
       "            [ 6.36858e-01,  1.61863e-01,  9.47738e-01,  ..., -6.99489e+00, -4.75545e+00, -3.43741e+00],\n",
       "            ...,\n",
       "            [ 1.06300e+00, -4.65959e-01,  5.62975e-01,  ..., -9.96357e+00, -5.87908e+00, -5.81003e+00],\n",
       "            [ 7.05979e-01, -6.34384e-02,  4.24215e-01,  ..., -1.09256e+01, -6.95265e+00, -6.86729e+00],\n",
       "            [ 2.95172e-01, -8.83697e-02,  3.48454e-01,  ..., -7.66372e+00, -5.73988e+00, -4.61694e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.64828e-01,  1.63969e-01,  6.76177e-01,  ..., -4.71849e+00, -3.16294e+00, -3.15511e+00],\n",
       "            [ 6.59857e-01, -9.82673e-03,  1.28410e+00,  ..., -5.50080e+00, -4.29665e+00, -3.53460e+00],\n",
       "            [ 5.15197e-01, -9.38664e-02,  1.11745e+00,  ..., -6.44667e+00, -4.41942e+00, -3.81218e+00],\n",
       "            ...,\n",
       "            [ 1.21880e-01,  6.79605e-01,  5.43512e-01,  ..., -8.09097e+00, -5.43438e+00, -4.30376e+00],\n",
       "            [ 7.66990e-01,  4.93302e-01,  5.41016e-01,  ..., -8.23601e+00, -5.37181e+00, -3.88541e+00],\n",
       "            [ 3.38329e-01,  6.27826e-01,  3.38509e-01,  ..., -5.69982e+00, -4.64872e+00, -2.37550e+00]],\n",
       " \n",
       "           [[-4.26816e-02,  1.07331e+00,  7.89964e-01,  ..., -4.89835e+00, -4.16909e+00, -2.25498e+00],\n",
       "            [ 3.85183e-01,  1.29741e+00,  1.37022e+00,  ..., -5.23291e+00, -4.36250e+00, -2.69675e+00],\n",
       "            [ 1.54946e-01,  1.22568e+00,  7.77264e-01,  ..., -6.96806e+00, -5.50778e+00, -3.41623e+00],\n",
       "            ...,\n",
       "            [ 1.71762e-01,  9.29963e-01,  4.16758e-01,  ..., -8.50659e+00, -6.28979e+00, -5.13025e+00],\n",
       "            [ 9.13903e-01,  1.05770e+00,  9.46294e-01,  ..., -7.77101e+00, -5.46880e+00, -3.23866e+00],\n",
       "            [ 2.89827e-01,  7.87081e-01,  1.41571e-01,  ..., -6.78349e+00, -5.67831e+00, -2.98133e+00]],\n",
       " \n",
       "           [[-2.81385e-02, -2.62423e-01,  7.44683e-01,  ..., -5.15145e+00, -5.23484e+00, -1.25023e+00],\n",
       "            [ 4.95441e-01, -8.70662e-02,  1.10525e+00,  ..., -5.47057e+00, -4.94478e+00, -1.95277e+00],\n",
       "            [ 4.12933e-01, -1.44311e-01,  7.00633e-01,  ..., -6.85800e+00, -5.64167e+00, -2.49982e+00],\n",
       "            ...,\n",
       "            [ 3.11199e-01, -3.29226e-01,  3.54672e-01,  ..., -7.55824e+00, -6.24633e+00, -3.88033e+00],\n",
       "            [ 8.03061e-01, -3.07527e-01,  4.29145e-01,  ..., -7.65323e+00, -5.93759e+00, -2.95116e+00],\n",
       "            [ 3.02349e-01, -5.38519e-01,  1.81686e-01,  ..., -6.28780e+00, -5.36096e+00, -2.91118e+00]]],\n",
       " \n",
       " \n",
       "          [[[-3.08949e-02,  4.76795e-01, -1.06666e-01,  ..., -4.42697e+00, -4.84581e+00, -1.92849e+00],\n",
       "            [-2.26530e-01,  1.61561e-01,  3.47693e-01,  ..., -5.14382e+00, -4.45690e+00, -1.61712e+00],\n",
       "            [ 1.15777e-01,  3.72598e-01,  1.04843e-01,  ..., -6.38548e+00, -4.44544e+00, -2.75897e+00],\n",
       "            ...,\n",
       "            [ 1.46973e-01,  3.68550e-01,  2.28444e-01,  ..., -6.13509e+00, -4.11355e+00, -3.33417e+00],\n",
       "            [-1.21446e-01,  3.46976e-01,  1.33764e-01,  ..., -6.90256e+00, -5.44940e+00, -3.76052e+00],\n",
       "            [-2.52330e-01,  2.41811e-01,  1.32273e-01,  ..., -4.17432e+00, -4.77378e+00, -1.92700e+00]],\n",
       " \n",
       "           [[-1.55616e-01,  2.18423e-01, -3.05104e-01,  ..., -5.59778e+00, -4.97459e+00, -2.66317e+00],\n",
       "            [-5.06979e-01, -1.04835e-01,  1.77017e-01,  ..., -7.06108e+00, -4.95703e+00, -4.05177e+00],\n",
       "            [-3.60973e-02, -1.57379e-01,  2.36025e-01,  ..., -6.89726e+00, -4.52727e+00, -3.53713e+00],\n",
       "            ...,\n",
       "            [ 5.62858e-02,  6.60455e-01,  5.46960e-03,  ..., -8.86874e+00, -5.30528e+00, -5.41194e+00],\n",
       "            [-8.14770e-02,  1.66666e-01,  1.59542e-01,  ..., -7.76828e+00, -5.43586e+00, -3.58809e+00],\n",
       "            [-4.05647e-02,  5.74879e-01, -6.26473e-02,  ..., -5.26749e+00, -5.14150e+00, -2.96918e+00]],\n",
       " \n",
       "           [[-2.48938e-01,  4.97066e-01, -2.91006e-01,  ..., -5.48302e+00, -4.72610e+00, -2.48806e+00],\n",
       "            [-7.40098e-01, -1.50491e-01,  4.22750e-01,  ..., -6.31422e+00, -4.24105e+00, -3.72033e+00],\n",
       "            [-5.24185e-02,  5.26025e-01,  4.67119e-01,  ..., -6.38149e+00, -3.94364e+00, -3.78380e+00],\n",
       "            ...,\n",
       "            [ 2.29330e-01, -1.73830e-01,  2.72504e-02,  ..., -9.17835e+00, -4.90631e+00, -5.89473e+00],\n",
       "            [ 2.55754e-01, -5.45044e-02, -1.74985e-01,  ..., -1.02381e+01, -5.99794e+00, -6.79265e+00],\n",
       "            [ 1.60982e-01,  2.10924e-01, -2.12368e-01,  ..., -6.85444e+00, -5.31023e+00, -4.39996e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 5.65943e-03,  6.35063e-01,  9.47689e-02,  ..., -3.51696e+00, -3.32287e+00, -3.24760e+00],\n",
       "            [ 1.67731e-01,  2.63961e-01,  6.56870e-01,  ..., -4.87984e+00, -3.62816e+00, -3.73921e+00],\n",
       "            [ 1.22510e-01,  2.19811e-01,  5.86543e-01,  ..., -5.81599e+00, -3.61697e+00, -4.31952e+00],\n",
       "            ...,\n",
       "            [-5.66655e-01,  9.31515e-01,  9.57126e-02,  ..., -7.38930e+00, -4.43638e+00, -5.01023e+00],\n",
       "            [-1.04843e-01,  7.99095e-01,  1.58894e-01,  ..., -7.58492e+00, -4.52589e+00, -4.54393e+00],\n",
       "            [ 8.75754e-03,  1.13430e+00, -1.16718e-01,  ..., -4.86041e+00, -4.25894e+00, -2.48587e+00]],\n",
       " \n",
       "           [[ 9.38965e-02,  1.32265e+00,  4.23893e-02,  ..., -3.79926e+00, -4.20408e+00, -2.47742e+00],\n",
       "            [-3.82344e-01,  1.57940e+00,  6.42019e-01,  ..., -4.33369e+00, -3.75357e+00, -3.23531e+00],\n",
       "            [-2.07997e-01,  1.48042e+00,  2.10816e-01,  ..., -6.08965e+00, -4.75826e+00, -4.04354e+00],\n",
       "            ...,\n",
       "            [-1.74761e-01,  1.05096e+00, -1.52701e-01,  ..., -7.45615e+00, -5.58757e+00, -5.58126e+00],\n",
       "            [-1.29775e-01,  1.14174e+00,  4.17776e-01,  ..., -6.86776e+00, -4.50453e+00, -4.26490e+00],\n",
       "            [ 2.15139e-01,  1.11117e+00, -3.57160e-01,  ..., -6.20069e+00, -5.34195e+00, -2.78220e+00]],\n",
       " \n",
       "           [[ 6.22943e-02, -3.21262e-02,  2.79956e-02,  ..., -4.45001e+00, -5.14791e+00, -1.10892e+00],\n",
       "            [-2.03594e-01, -7.05951e-02,  5.01556e-01,  ..., -4.66632e+00, -4.43579e+00, -1.88940e+00],\n",
       "            [-1.47934e-01, -6.86931e-02,  2.05350e-01,  ..., -6.07782e+00, -5.19550e+00, -2.55451e+00],\n",
       "            ...,\n",
       "            [-1.36463e-01, -1.08685e-01, -1.85231e-01,  ..., -6.68133e+00, -5.85135e+00, -4.00374e+00],\n",
       "            [-8.04268e-02, -2.95341e-01, -7.40191e-02,  ..., -6.74904e+00, -5.34603e+00, -3.44150e+00],\n",
       "            [ 1.14907e-01, -3.22599e-01, -3.29023e-01,  ..., -5.56258e+00, -5.28025e+00, -2.78334e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-3.52729e-01, -2.04228e-01,  1.90275e-01,  ..., -6.97586e+00, -3.31311e+00, -6.15873e+00],\n",
       "            [ 2.58070e-01, -1.55309e-01,  8.74303e-01,  ..., -7.08713e+00, -3.64969e+00, -5.55332e+00],\n",
       "            [-6.92424e-01, -1.67813e-01,  4.34214e-01,  ..., -7.40268e+00, -4.62837e+00, -5.55016e+00],\n",
       "            ...,\n",
       "            [-9.09745e-02, -3.33569e-01,  4.00713e-01,  ..., -7.99231e+00, -4.22137e+00, -6.10959e+00],\n",
       "            [ 4.19026e-02, -3.88124e-01,  6.54405e-01,  ..., -8.14708e+00, -3.71178e+00, -5.83420e+00],\n",
       "            [ 1.84525e-01,  1.01268e-04,  5.65158e-01,  ..., -6.60819e+00, -2.32062e+00, -5.89272e+00]],\n",
       " \n",
       "           [[-4.86282e-01, -7.40379e-03, -4.22733e-02,  ..., -8.68022e+00, -4.08842e+00, -6.15115e+00],\n",
       "            [-1.17870e-01,  1.12943e+00,  6.94981e-01,  ..., -9.03078e+00, -3.86466e+00, -5.52381e+00],\n",
       "            [ 1.17374e-01,  7.02729e-01,  7.26739e-01,  ..., -8.45631e+00, -4.41502e+00, -4.94346e+00],\n",
       "            ...,\n",
       "            [ 1.23986e-01, -4.40829e-01,  3.90988e-01,  ..., -8.81449e+00, -4.94663e+00, -5.61236e+00],\n",
       "            [ 7.09655e-01, -3.13112e-01,  5.17673e-01,  ..., -9.41888e+00, -4.72569e+00, -5.77212e+00],\n",
       "            [ 3.89936e-01, -2.29411e-01,  1.66371e-01,  ..., -7.49426e+00, -2.75817e+00, -5.85103e+00]],\n",
       " \n",
       "           [[-2.59693e-01,  9.09734e-01,  2.06754e-01,  ..., -8.58827e+00, -4.03594e+00, -6.17920e+00],\n",
       "            [ 2.37719e-01,  3.52740e-02,  1.47483e-01,  ..., -1.14610e+01, -5.92289e+00, -7.09497e+00],\n",
       "            [-7.19556e-01,  5.56991e-01,  3.36259e-01,  ..., -1.08006e+01, -6.10936e+00, -7.05220e+00],\n",
       "            ...,\n",
       "            [-2.12230e-01,  1.03806e+00,  4.05417e-01,  ..., -9.75829e+00, -5.28209e+00, -6.58702e+00],\n",
       "            [-2.48325e-01,  5.81601e-01,  3.38876e-01,  ..., -1.00750e+01, -5.89000e+00, -6.37159e+00],\n",
       "            [ 3.25526e-01, -3.98113e-01,  1.42121e-01,  ..., -7.56818e+00, -3.25495e+00, -5.95150e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.78635e-01, -3.60458e-02, -4.54755e-02,  ..., -7.71586e+00, -3.89263e+00, -5.78874e+00],\n",
       "            [-4.87774e-02,  6.36499e-01,  5.93741e-01,  ..., -7.99575e+00, -4.96614e+00, -5.73606e+00],\n",
       "            [-2.95039e-02, -7.49201e-01,  4.66746e-01,  ..., -8.82058e+00, -5.17602e+00, -5.74889e+00],\n",
       "            ...,\n",
       "            [-4.09670e-02, -2.90988e-01,  3.72450e-01,  ..., -9.50712e+00, -4.86659e+00, -5.29298e+00],\n",
       "            [-2.23855e-01, -4.27603e-01,  3.12900e-01,  ..., -1.02504e+01, -5.26980e+00, -5.71905e+00],\n",
       "            [ 4.34955e-01, -2.53394e-01,  7.27746e-02,  ..., -8.43173e+00, -3.72787e+00, -6.39481e+00]],\n",
       " \n",
       "           [[-5.88178e-01,  5.67300e-01,  2.19936e-02,  ..., -6.94239e+00, -4.07965e+00, -5.56357e+00],\n",
       "            [-1.25479e-01,  8.78208e-01,  8.17098e-01,  ..., -7.28949e+00, -4.43023e+00, -5.35867e+00],\n",
       "            [ 1.90214e-01, -1.13401e-01,  7.76294e-01,  ..., -7.92307e+00, -4.19211e+00, -5.11578e+00],\n",
       "            ...,\n",
       "            [ 1.31063e-01,  1.48842e+00,  6.27579e-01,  ..., -9.15186e+00, -3.73022e+00, -6.13950e+00],\n",
       "            [-4.10189e-03,  1.51927e+00,  7.04665e-01,  ..., -9.62076e+00, -4.52963e+00, -6.30752e+00],\n",
       "            [ 2.54570e-01,  8.85856e-01,  2.94700e-01,  ..., -7.35497e+00, -2.30579e+00, -5.61208e+00]],\n",
       " \n",
       "           [[-1.70992e-01, -6.47458e-02,  3.38862e-01,  ..., -6.40031e+00, -4.76236e+00, -6.46363e+00],\n",
       "            [-4.93625e-01,  1.23272e-01,  5.89025e-01,  ..., -7.36329e+00, -4.65413e+00, -6.28821e+00],\n",
       "            [ 2.31876e-01,  2.03640e-01,  5.82969e-01,  ..., -7.93613e+00, -4.49307e+00, -6.26818e+00],\n",
       "            ...,\n",
       "            [ 1.02368e-01,  1.92261e-01,  3.99113e-01,  ..., -8.75192e+00, -3.72370e+00, -6.63050e+00],\n",
       "            [ 4.63861e-01, -8.10392e-02,  4.04929e-01,  ..., -8.35412e+00, -3.49096e+00, -6.39032e+00],\n",
       "            [ 4.09610e-01, -1.29877e-01,  5.94905e-02,  ..., -7.34278e+00, -3.16672e+00, -6.63553e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.28953e-01, -3.11444e-01, -2.97700e-01,  ..., -5.66755e+00, -2.38587e+00, -3.19085e+00],\n",
       "            [ 3.40609e-01, -2.82561e-01,  3.39075e-01,  ..., -5.01452e+00, -2.95967e+00, -4.29253e+00],\n",
       "            [-6.17768e-01, -2.59942e-01,  1.89880e-01,  ..., -5.09801e+00, -3.40808e+00, -3.84050e+00],\n",
       "            ...,\n",
       "            [-7.20996e-02, -4.19752e-01,  6.20183e-02,  ..., -5.38012e+00, -2.76008e+00, -4.11024e+00],\n",
       "            [ 1.15762e-01, -4.99163e-01,  2.69548e-01,  ..., -5.30016e+00, -2.68683e+00, -4.04391e+00],\n",
       "            [ 5.87277e-02, -1.13800e-01, -8.56142e-02,  ..., -5.52232e+00, -2.25045e+00, -3.87331e+00]],\n",
       " \n",
       "           [[-4.06550e-01, -1.03801e-01, -4.56320e-01,  ..., -5.46657e+00, -2.80014e+00, -4.40195e+00],\n",
       "            [-2.15577e-01,  8.03534e-01,  1.85024e-01,  ..., -4.61249e+00, -2.49213e+00, -6.72017e+00],\n",
       "            [ 9.38675e-02,  3.20184e-01,  1.79085e-01,  ..., -4.52591e+00, -2.66989e+00, -5.95454e+00],\n",
       "            ...,\n",
       "            [ 1.24088e-01, -5.49438e-01,  6.15039e-02,  ..., -5.06274e+00, -3.46378e+00, -5.91285e+00],\n",
       "            [ 7.64627e-01, -4.07800e-01,  1.00121e-01,  ..., -4.99225e+00, -3.09643e+00, -6.05802e+00],\n",
       "            [ 3.36189e-01, -3.45508e-01, -3.55399e-01,  ..., -5.43559e+00, -2.46809e+00, -4.49408e+00]],\n",
       " \n",
       "           [[-2.24583e-01,  6.92266e-01, -2.81196e-01,  ..., -5.36101e+00, -3.15437e+00, -5.19803e+00],\n",
       "            [ 2.82878e-01, -1.30210e-02, -2.91471e-01,  ..., -5.76469e+00, -3.40487e+00, -5.49970e+00],\n",
       "            [-7.69065e-01,  4.79895e-01, -1.86047e-01,  ..., -5.81774e+00, -4.07548e+00, -4.74959e+00],\n",
       "            ...,\n",
       "            [-2.40160e-01,  9.12994e-01, -7.84466e-02,  ..., -5.45629e+00, -3.48021e+00, -5.74612e+00],\n",
       "            [-2.94930e-01,  4.99832e-01, -9.29990e-02,  ..., -5.30062e+00, -3.47752e+00, -5.66074e+00],\n",
       "            [ 2.92160e-01, -5.06099e-01, -3.35173e-01,  ..., -5.53148e+00, -2.60390e+00, -4.86830e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.93574e-01, -8.31726e-02, -4.69054e-01,  ..., -5.09003e+00, -3.21091e+00, -4.62292e+00],\n",
       "            [-1.62743e-01,  6.81088e-01,  2.26388e-01,  ..., -4.60990e+00, -4.07572e+00, -6.31350e+00],\n",
       "            [-8.73085e-02, -7.27341e-01, -1.79027e-02,  ..., -4.96854e+00, -3.82318e+00, -5.98688e+00],\n",
       "            ...,\n",
       "            [-1.13855e-01, -3.58745e-01,  6.67495e-03,  ..., -4.67326e+00, -2.65043e+00, -7.51207e+00],\n",
       "            [-2.70392e-01, -4.65574e-01,  1.29812e-02,  ..., -4.79733e+00, -3.05001e+00, -7.88622e+00],\n",
       "            [ 4.18359e-01, -3.08662e-01, -4.35649e-01,  ..., -5.54889e+00, -2.93510e+00, -5.47810e+00]],\n",
       " \n",
       "           [[-4.70072e-01,  5.18668e-01, -4.07803e-01,  ..., -5.03530e+00, -3.64240e+00, -3.72914e+00],\n",
       "            [-1.72603e-01,  8.13738e-01,  3.00182e-01,  ..., -4.46233e+00, -4.08535e+00, -5.11524e+00],\n",
       "            [ 5.53432e-02, -2.01378e-01,  2.58847e-01,  ..., -4.62714e+00, -3.67583e+00, -5.76471e+00],\n",
       "            ...,\n",
       "            [ 1.43447e-01,  1.34696e+00,  6.17802e-02,  ..., -4.81920e+00, -2.63599e+00, -7.12543e+00],\n",
       "            [ 2.95140e-02,  1.41089e+00,  1.84333e-01,  ..., -4.95390e+00, -3.58330e+00, -7.61832e+00],\n",
       "            [ 1.69451e-01,  7.05568e-01, -2.39138e-01,  ..., -4.98204e+00, -2.73957e+00, -6.19456e+00]],\n",
       " \n",
       "           [[ 1.51536e-01, -7.70818e-03,  1.68373e-02,  ..., -5.69633e+00, -4.09692e+00, -2.24200e+00],\n",
       "            [-4.82605e-01,  8.98565e-02,  4.22083e-01,  ..., -5.23447e+00, -3.98196e+00, -3.63283e+00],\n",
       "            [ 2.21568e-01,  1.33640e-01,  3.10367e-01,  ..., -5.28960e+00, -3.51661e+00, -4.01456e+00],\n",
       "            ...,\n",
       "            [ 1.85590e-01,  1.32090e-01, -1.46371e-02,  ..., -5.48777e+00, -2.94741e+00, -4.68078e+00],\n",
       "            [ 5.97102e-01, -1.55043e-01,  7.29217e-02,  ..., -5.35432e+00, -3.37179e+00, -4.62181e+00],\n",
       "            [ 3.27753e-01, -1.47449e-01, -4.44609e-01,  ..., -5.80870e+00, -3.18179e+00, -3.68415e+00]]],\n",
       " \n",
       " \n",
       "          [[[-2.17190e-01, -1.32181e-01, -9.93412e-02,  ..., -5.72652e+00, -2.14003e+00, -3.98952e+00],\n",
       "            [ 3.45817e-01, -9.98715e-02,  6.95160e-01,  ..., -5.75268e+00, -2.79438e+00, -4.30284e+00],\n",
       "            [-5.87944e-01, -1.51539e-01,  5.09481e-01,  ..., -6.01300e+00, -3.84253e+00, -4.06646e+00],\n",
       "            ...,\n",
       "            [-6.12640e-02, -3.07031e-01,  3.54092e-01,  ..., -6.18555e+00, -2.98398e+00, -4.44450e+00],\n",
       "            [ 1.24749e-01, -3.39662e-01,  6.27114e-01,  ..., -6.27267e+00, -2.41703e+00, -4.58339e+00],\n",
       "            [ 5.02373e-02,  7.51385e-02,  1.80704e-01,  ..., -5.57495e+00, -1.18164e+00, -4.40382e+00]],\n",
       " \n",
       "           [[-3.86517e-01, -6.98476e-02, -2.84496e-01,  ..., -6.94483e+00, -2.84105e+00, -4.93080e+00],\n",
       "            [-1.89781e-01,  8.70716e-01,  4.88955e-01,  ..., -7.43344e+00, -2.53989e+00, -7.11762e+00],\n",
       "            [ 1.27890e-01,  3.97098e-01,  5.26551e-01,  ..., -6.94439e+00, -2.97013e+00, -6.90453e+00],\n",
       "            ...,\n",
       "            [ 1.49792e-01, -5.31390e-01,  3.53077e-01,  ..., -7.03169e+00, -3.70168e+00, -5.90362e+00],\n",
       "            [ 7.81091e-01, -3.86353e-01,  4.08595e-01,  ..., -7.53038e+00, -3.45732e+00, -6.14798e+00],\n",
       "            [ 3.33779e-01, -3.22163e-01, -1.00890e-01,  ..., -6.01641e+00, -1.40142e+00, -4.82026e+00]],\n",
       " \n",
       "           [[-2.07616e-01,  7.72474e-01, -1.12551e-01,  ..., -6.64534e+00, -2.63494e+00, -6.04529e+00],\n",
       "            [ 3.04493e-01, -1.86500e-02, -8.29957e-02,  ..., -8.89437e+00, -4.14719e+00, -7.73894e+00],\n",
       "            [-7.37650e-01,  5.31478e-01,  3.49668e-02,  ..., -8.09385e+00, -4.63086e+00, -6.87183e+00],\n",
       "            ...,\n",
       "            [-2.07958e-01,  9.85754e-01,  1.42471e-01,  ..., -7.79643e+00, -3.67855e+00, -6.83727e+00],\n",
       "            [-2.71300e-01,  5.59860e-01,  1.45266e-01,  ..., -8.07547e+00, -4.40298e+00, -6.65937e+00],\n",
       "            [ 2.92741e-01, -4.29598e-01, -9.06347e-02,  ..., -6.48888e+00, -2.08503e+00, -5.27950e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.74228e-01, -3.94975e-02, -3.03288e-01,  ..., -6.17722e+00, -2.91377e+00, -4.76784e+00],\n",
       "            [-1.35284e-01,  7.39835e-01,  5.29787e-01,  ..., -6.97148e+00, -4.55357e+00, -5.06477e+00],\n",
       "            [-6.09773e-02, -7.49483e-01,  2.43687e-01,  ..., -7.32158e+00, -4.28061e+00, -6.16674e+00],\n",
       "            ...,\n",
       "            [-8.76986e-02, -3.96954e-01,  3.04203e-01,  ..., -8.18795e+00, -3.23313e+00, -9.00815e+00],\n",
       "            [-2.43816e-01, -5.01636e-01,  3.02672e-01,  ..., -8.56207e+00, -3.96709e+00, -7.98472e+00],\n",
       "            [ 4.16832e-01, -2.92581e-01, -2.27410e-01,  ..., -6.86604e+00, -2.53689e+00, -6.09950e+00]],\n",
       " \n",
       "           [[-4.48501e-01,  5.48458e-01, -2.51406e-01,  ..., -5.35852e+00, -3.12112e+00, -3.97340e+00],\n",
       "            [-1.64623e-01,  8.79241e-01,  6.06193e-01,  ..., -5.99171e+00, -4.07714e+00, -4.86579e+00],\n",
       "            [ 7.47596e-02, -1.41590e-01,  5.95806e-01,  ..., -6.53228e+00, -3.64827e+00, -5.69684e+00],\n",
       "            ...,\n",
       "            [ 1.67551e-01,  1.45253e+00,  2.76690e-01,  ..., -7.98391e+00, -2.96124e+00, -7.95744e+00],\n",
       "            [ 3.57338e-02,  1.53345e+00,  4.39795e-01,  ..., -8.06596e+00, -3.69550e+00, -7.03895e+00],\n",
       "            [ 1.60934e-01,  8.10326e-01,  1.01773e-02,  ..., -6.51997e+00, -1.57999e+00, -5.61879e+00]],\n",
       " \n",
       "           [[ 1.59081e-01, -1.23599e-02,  2.14768e-01,  ..., -4.67993e+00, -3.95372e+00, -1.74268e+00],\n",
       "            [-4.72406e-01,  9.22272e-02,  7.70375e-01,  ..., -5.56751e+00, -4.16643e+00, -2.56984e+00],\n",
       "            [ 2.30594e-01,  1.48058e-01,  6.25784e-01,  ..., -6.00161e+00, -3.70334e+00, -3.68286e+00],\n",
       "            ...,\n",
       "            [ 2.00957e-01,  1.66971e-01,  2.27937e-01,  ..., -6.75476e+00, -2.70747e+00, -4.09998e+00],\n",
       "            [ 6.00387e-01, -1.46236e-01,  4.11198e-01,  ..., -6.46169e+00, -2.59880e+00, -4.17354e+00],\n",
       "            [ 3.14587e-01, -1.69751e-01, -2.46452e-01,  ..., -5.74230e+00, -2.53100e+00, -2.93561e+00]]]]], device='cuda:0'),\n",
       " tensor([[[[[-4.52741e-01, -4.14671e-02, -5.91435e-01,  ..., -5.62737e+00, -4.24842e+00, -3.50413e+00],\n",
       "            [-4.95543e-01,  3.12076e-01, -2.66672e-01,  ..., -5.71931e+00, -5.83872e+00, -2.28965e+00],\n",
       "            [-9.02449e-02, -5.65752e-01, -1.64024e-01,  ..., -5.67140e+00, -6.51515e+00, -2.65939e+00],\n",
       "            ...,\n",
       "            [-2.88781e-01, -8.89322e-02, -2.36313e-01,  ..., -5.57944e+00, -6.04703e+00, -3.04939e+00],\n",
       "            [-8.68164e-02, -4.78436e-01, -9.95925e-02,  ..., -5.30122e+00, -4.96809e+00, -3.09045e+00],\n",
       "            [ 4.29260e-02, -5.96461e-02, -3.37531e-01,  ..., -5.12825e+00, -3.05739e+00, -2.88604e+00]],\n",
       " \n",
       "           [[-6.85536e-01,  2.89596e-01, -8.46161e-01,  ..., -5.78475e+00, -4.95742e+00, -2.89448e+00],\n",
       "            [-5.41056e-01, -2.66404e-02, -7.91704e-01,  ..., -6.18839e+00, -6.46698e+00, -2.97523e+00],\n",
       "            [ 3.38787e-01, -8.01558e-01, -5.82370e-01,  ..., -6.47139e+00, -7.95357e+00, -3.39206e+00],\n",
       "            ...,\n",
       "            [ 1.61078e-01,  1.18233e-01, -5.73456e-01,  ..., -6.60290e+00, -8.29431e+00, -3.91724e+00],\n",
       "            [ 3.95663e-01,  5.74929e-02, -5.99476e-01,  ..., -6.52556e+00, -8.00306e+00, -3.24298e+00],\n",
       "            [ 1.62492e-02, -4.54711e-01, -7.52134e-01,  ..., -6.05312e+00, -5.64607e+00, -2.65447e+00]],\n",
       " \n",
       "           [[ 5.77240e-01,  1.27808e-01, -7.32327e-01,  ..., -5.84739e+00, -5.56762e+00, -2.42029e+00],\n",
       "            [ 1.34385e-01, -4.96992e-01, -7.74016e-01,  ..., -6.53986e+00, -7.62192e+00, -3.21728e+00],\n",
       "            [ 1.91240e-01,  1.04318e-01, -7.88821e-01,  ..., -6.74604e+00, -7.89224e+00, -3.69873e+00],\n",
       "            ...,\n",
       "            [-2.73797e-01,  2.87878e-01, -1.64442e-01,  ..., -6.28662e+00, -7.64773e+00, -4.03442e+00],\n",
       "            [ 7.49643e-01,  2.72895e-02, -2.81449e-01,  ..., -6.40862e+00, -7.02522e+00, -3.71413e+00],\n",
       "            [ 2.15579e-01, -2.73789e-01, -6.46599e-01,  ..., -5.71115e+00, -4.84143e+00, -3.03052e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-8.83723e-01,  2.87115e-02, -7.96383e-01,  ..., -5.47236e+00, -3.69694e+00, -2.65217e+00],\n",
       "            [ 1.69012e-01, -5.28461e-01, -7.15849e-01,  ..., -6.22130e+00, -7.07364e+00, -3.13679e+00],\n",
       "            [-3.57886e-01, -6.88186e-01, -6.20818e-01,  ..., -6.25112e+00, -7.17522e+00, -3.16205e+00],\n",
       "            ...,\n",
       "            [-5.70698e-01, -9.11328e-01, -5.98967e-01,  ..., -6.09972e+00, -6.63824e+00, -3.41508e+00],\n",
       "            [ 1.93582e-01,  3.84465e-01, -5.42190e-01,  ..., -6.08154e+00, -5.96101e+00, -4.45606e+00],\n",
       "            [ 1.91616e-01,  5.34162e-01, -5.15719e-01,  ..., -5.12371e+00, -3.56779e+00, -3.65777e+00]],\n",
       " \n",
       "           [[-7.92756e-01,  4.53855e-01, -6.86597e-01,  ..., -5.60700e+00, -4.07897e+00, -3.06453e+00],\n",
       "            [-2.20965e-01,  4.88167e-01, -2.66147e-01,  ..., -5.97049e+00, -6.68768e+00, -3.58610e+00],\n",
       "            [ 4.91396e-01,  6.74594e-01, -2.33531e-01,  ..., -5.72306e+00, -6.42607e+00, -3.47959e+00],\n",
       "            ...,\n",
       "            [ 4.72735e-01,  2.27356e-01, -6.12092e-01,  ..., -6.26059e+00, -7.44355e+00, -4.39343e+00],\n",
       "            [ 5.05438e-02, -1.36115e+00, -5.88272e-01,  ..., -5.76939e+00, -5.29721e+00, -4.99663e+00],\n",
       "            [-5.16168e-01,  3.32225e-01, -3.63594e-01,  ..., -5.09833e+00, -4.81577e+00, -4.11145e+00]],\n",
       " \n",
       "           [[-5.00116e-02,  2.85565e-01, -1.67010e-01,  ..., -4.78911e+00, -4.36149e+00, -2.25695e+00],\n",
       "            [-7.39205e-01,  2.43306e-02,  1.24413e-01,  ..., -4.67351e+00, -5.22868e+00, -2.12069e+00],\n",
       "            [-4.53926e-02, -3.23347e-01, -1.63438e-01,  ..., -5.02654e+00, -5.96703e+00, -2.27764e+00],\n",
       "            ...,\n",
       "            [-2.62851e-03,  1.91323e-01, -3.70263e-01,  ..., -5.18723e+00, -5.00565e+00, -3.10246e+00],\n",
       "            [ 6.75140e-01,  4.49111e-02, -2.07418e-01,  ..., -5.06101e+00, -5.18429e+00, -3.24672e+00],\n",
       "            [-2.25657e-01,  1.40555e-01, -2.41614e-01,  ..., -4.79289e+00, -4.19027e+00, -2.79043e+00]]],\n",
       " \n",
       " \n",
       "          [[[-4.66104e-01, -2.54638e-02, -5.12880e-01,  ..., -5.32927e+00, -5.31489e+00, -7.00899e+00],\n",
       "            [-4.79088e-01,  3.27527e-01, -2.64665e-01,  ..., -7.49399e+00, -6.90484e+00, -6.89057e+00],\n",
       "            [-7.96545e-02, -5.49257e-01, -1.02255e-01,  ..., -5.83597e+00, -6.75884e+00, -7.02145e+00],\n",
       "            ...,\n",
       "            [-2.79393e-01, -4.19111e-02, -2.50942e-01,  ..., -6.98178e+00, -6.78096e+00, -7.75409e+00],\n",
       "            [-7.36690e-02, -4.50206e-01, -3.71555e-02,  ..., -5.49768e+00, -5.87255e+00, -7.36907e+00],\n",
       "            [ 5.45514e-02, -4.50050e-02, -2.91783e-01,  ..., -3.79821e+00, -4.08699e+00, -6.42178e+00]],\n",
       " \n",
       "           [[-6.95347e-01,  2.63116e-01, -7.94315e-01,  ..., -7.23405e+00, -6.13898e+00, -7.53119e+00],\n",
       "            [-5.40604e-01, -4.51334e-02, -8.41806e-01,  ..., -1.04826e+01, -7.83559e+00, -8.27653e+00],\n",
       "            [ 3.46853e-01, -8.54576e-01, -6.47174e-01,  ..., -1.18123e+01, -9.30875e+00, -8.37272e+00],\n",
       "            ...,\n",
       "            [ 1.75116e-01,  1.47932e-01, -6.59596e-01,  ..., -1.07616e+01, -8.98833e+00, -8.86933e+00],\n",
       "            [ 4.11949e-01,  9.84310e-02, -6.47776e-01,  ..., -8.78305e+00, -8.33612e+00, -8.18738e+00],\n",
       "            [ 1.94275e-02, -4.27247e-01, -7.47924e-01,  ..., -6.38869e+00, -6.99724e+00, -8.04797e+00]],\n",
       " \n",
       "           [[ 5.81223e-01,  1.11598e-01, -7.19028e-01,  ..., -7.12394e+00, -6.51331e+00, -7.27694e+00],\n",
       "            [ 1.29579e-01, -5.18964e-01, -8.09384e-01,  ..., -1.01276e+01, -8.41794e+00, -8.46176e+00],\n",
       "            [ 1.96749e-01,  7.91333e-02, -8.21730e-01,  ..., -1.24985e+01, -9.15849e+00, -8.68235e+00],\n",
       "            ...,\n",
       "            [-2.69309e-01,  3.30496e-01, -1.62685e-01,  ..., -1.07923e+01, -8.64952e+00, -8.00571e+00],\n",
       "            [ 7.70772e-01,  8.10389e-02, -2.73902e-01,  ..., -1.00585e+01, -8.31078e+00, -8.40907e+00],\n",
       "            [ 2.30694e-01, -2.45528e-01, -6.67015e-01,  ..., -7.28255e+00, -6.72262e+00, -8.07937e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-8.94380e-01,  5.37833e-02, -7.78848e-01,  ..., -5.72595e+00, -4.50099e+00, -6.98193e+00],\n",
       "            [ 1.64337e-01, -4.96066e-01, -7.92640e-01,  ..., -9.95762e+00, -8.18949e+00, -8.36009e+00],\n",
       "            [-3.58015e-01, -6.52232e-01, -7.16903e-01,  ..., -1.09155e+01, -8.02419e+00, -8.19119e+00],\n",
       "            ...,\n",
       "            [-5.73135e-01, -8.95422e-01, -6.97667e-01,  ..., -1.05359e+01, -7.40144e+00, -8.34067e+00],\n",
       "            [ 2.01658e-01,  4.32996e-01, -6.37279e-01,  ..., -1.15945e+01, -8.25199e+00, -8.97337e+00],\n",
       "            [ 2.01655e-01,  5.50723e-01, -5.69485e-01,  ..., -7.90071e+00, -5.85415e+00, -7.82713e+00]],\n",
       " \n",
       "           [[-8.09521e-01,  4.70049e-01, -6.29599e-01,  ..., -5.70872e+00, -4.76619e+00, -6.49009e+00],\n",
       "            [-2.18569e-01,  5.12985e-01, -2.34789e-01,  ..., -9.26782e+00, -8.37694e+00, -7.61029e+00],\n",
       "            [ 5.02405e-01,  7.10104e-01, -2.50733e-01,  ..., -9.13343e+00, -7.01967e+00, -7.33026e+00],\n",
       "            ...,\n",
       "            [ 4.87990e-01,  2.74516e-01, -7.01877e-01,  ..., -1.13098e+01, -8.58040e+00, -8.75608e+00],\n",
       "            [ 5.09772e-02, -1.34596e+00, -6.17842e-01,  ..., -1.14488e+01, -7.64953e+00, -8.73636e+00],\n",
       "            [-5.04363e-01,  3.50753e-01, -3.74830e-01,  ..., -7.69405e+00, -6.51379e+00, -7.61013e+00]],\n",
       " \n",
       "           [[-5.57594e-02,  2.71057e-01, -5.03180e-02,  ..., -3.69483e+00, -4.71798e+00, -4.87265e+00],\n",
       "            [-7.24522e-01, -3.19939e-03,  2.04735e-01,  ..., -5.08738e+00, -5.88639e+00, -5.52542e+00],\n",
       "            [-3.81274e-02, -3.49604e-01, -1.57519e-01,  ..., -6.99845e+00, -6.83099e+00, -6.42560e+00],\n",
       "            ...,\n",
       "            [ 4.86115e-03,  1.97877e-01, -4.22787e-01,  ..., -8.11867e+00, -6.16312e+00, -6.82600e+00],\n",
       "            [ 6.87044e-01,  5.73217e-02, -2.03054e-01,  ..., -6.53256e+00, -6.27129e+00, -6.78610e+00],\n",
       "            [-2.17992e-01,  1.30427e-01, -1.66290e-01,  ..., -4.18703e+00, -4.93103e+00, -5.96023e+00]]],\n",
       " \n",
       " \n",
       "          [[[-3.21302e-01, -1.13077e-01, -7.40115e-01,  ..., -4.98687e+00, -3.39285e+00, -5.07348e+00],\n",
       "            [-4.52755e-01,  2.23429e-01, -3.66335e-01,  ..., -6.13935e+00, -4.07631e+00, -5.44100e+00],\n",
       "            [-6.91734e-02, -5.74297e-01, -3.43238e-01,  ..., -5.20074e+00, -4.63491e+00, -5.33063e+00],\n",
       "            ...,\n",
       "            [-2.42614e-01, -9.04456e-02, -2.88064e-01,  ..., -5.79802e+00, -4.00026e+00, -6.20094e+00],\n",
       "            [-7.49152e-02, -4.56016e-01, -2.36286e-01,  ..., -4.92157e+00, -3.53970e+00, -5.63939e+00],\n",
       "            [-1.02952e-02, -8.99799e-02, -5.52056e-01,  ..., -4.07091e+00, -2.42713e+00, -4.23179e+00]],\n",
       " \n",
       "           [[-5.49253e-01,  1.52123e-01, -8.94301e-01,  ..., -5.40958e+00, -3.69816e+00, -5.26037e+00],\n",
       "            [-4.73553e-01, -1.23062e-01, -7.71428e-01,  ..., -6.86330e+00, -4.16649e+00, -6.94651e+00],\n",
       "            [ 3.42910e-01, -9.06188e-01, -6.25292e-01,  ..., -7.71238e+00, -5.37860e+00, -7.28497e+00],\n",
       "            ...,\n",
       "            [ 1.88538e-01,  8.71775e-02, -4.94312e-01,  ..., -7.93286e+00, -5.51148e+00, -7.44875e+00],\n",
       "            [ 3.76572e-01,  7.85814e-02, -5.90669e-01,  ..., -6.94505e+00, -5.52976e+00, -6.15727e+00],\n",
       "            [-2.17452e-02, -5.54545e-01, -7.64201e-01,  ..., -4.96414e+00, -3.55301e+00, -5.57378e+00]],\n",
       " \n",
       "           [[ 6.66133e-01,  2.68509e-02, -7.66872e-01,  ..., -5.56196e+00, -4.00868e+00, -5.16628e+00],\n",
       "            [ 1.95134e-01, -6.04658e-01, -8.16844e-01,  ..., -6.84511e+00, -5.27818e+00, -6.77380e+00],\n",
       "            [ 2.13850e-01,  3.72510e-02, -7.70612e-01,  ..., -8.07563e+00, -5.64684e+00, -7.52895e+00],\n",
       "            ...,\n",
       "            [-1.93048e-01,  1.89765e-01, -1.76421e-01,  ..., -8.57520e+00, -4.86997e+00, -7.20589e+00],\n",
       "            [ 7.36208e-01, -7.71074e-03, -2.18272e-01,  ..., -7.45245e+00, -4.48919e+00, -6.99932e+00],\n",
       "            [ 1.30670e-01, -3.25038e-01, -7.26975e-01,  ..., -5.47131e+00, -3.60692e+00, -6.24818e+00]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-7.32450e-01, -2.22330e-02, -8.12347e-01,  ..., -5.25531e+00, -2.89321e+00, -4.55954e+00],\n",
       "            [ 2.39752e-01, -6.08521e-01, -6.81538e-01,  ..., -7.12881e+00, -4.76730e+00, -6.85893e+00],\n",
       "            [-2.74528e-01, -7.77795e-01, -6.31171e-01,  ..., -7.60609e+00, -4.55191e+00, -7.08647e+00],\n",
       "            ...,\n",
       "            [-5.43405e-01, -9.50493e-01, -7.32932e-01,  ..., -7.57907e+00, -4.49659e+00, -7.26775e+00],\n",
       "            [ 1.96658e-01,  3.34222e-01, -5.73122e-01,  ..., -7.49809e+00, -3.48984e+00, -7.95508e+00],\n",
       "            [ 1.07551e-01,  4.44990e-01, -6.02891e-01,  ..., -6.04957e+00, -2.46080e+00, -6.67718e+00]],\n",
       " \n",
       "           [[-6.48671e-01,  3.68523e-01, -8.20715e-01,  ..., -5.11912e+00, -3.36587e+00, -4.22117e+00],\n",
       "            [-1.40567e-01,  4.57456e-01, -3.33409e-01,  ..., -6.54746e+00, -4.92137e+00, -5.65023e+00],\n",
       "            [ 5.66594e-01,  6.39146e-01, -2.67120e-01,  ..., -6.77838e+00, -4.18896e+00, -6.12914e+00],\n",
       "            ...,\n",
       "            [ 4.77715e-01,  2.25638e-01, -5.97622e-01,  ..., -7.56393e+00, -4.74884e+00, -7.51402e+00],\n",
       "            [ 8.50564e-02, -1.42934e+00, -6.50753e-01,  ..., -6.91923e+00, -3.37867e+00, -8.12116e+00],\n",
       "            [-5.76542e-01,  3.14497e-01, -4.87910e-01,  ..., -5.50585e+00, -3.23773e+00, -6.24108e+00]],\n",
       " \n",
       "           [[ 1.29511e-02,  2.16465e-01, -5.14894e-01,  ..., -4.46855e+00, -4.20406e+00, -3.04905e+00],\n",
       "            [-6.81691e-01, -4.11028e-02, -1.48336e-01,  ..., -5.26427e+00, -4.27359e+00, -4.06294e+00],\n",
       "            [-1.13329e-03, -4.05112e-01, -3.18769e-01,  ..., -5.69750e+00, -4.67304e+00, -4.96151e+00],\n",
       "            ...,\n",
       "            [ 1.66629e-02,  1.00923e-01, -4.67631e-01,  ..., -6.21010e+00, -3.73802e+00, -5.74877e+00],\n",
       "            [ 6.89756e-01, -7.43462e-03, -3.82785e-01,  ..., -5.32684e+00, -3.83489e+00, -5.40905e+00],\n",
       "            [-2.47054e-01,  9.36951e-02, -6.08444e-01,  ..., -4.20494e+00, -3.95178e+00, -3.74916e+00]]]]], device='cuda:0'))"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trt_mod(input_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2b3366fc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fusing layers... \n",
      "IDetect.fuse\n"
     ]
    }
   ],
   "source": [
    "from test import pseudo_inference\n",
    "ci = pseudo_inference(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "371ad909",
   "metadata": {},
   "outputs": [],
   "source": [
    "B,C,H,W = input_image.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "13d22a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_model.eval().cpu()\n",
    "t1, t2 = torch_model(input_image.cpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "89aaccd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1, o2 = ci.run(trt_mod(input_image), nx_list=[W//8, W//16, W//32], ny_list=[H//8, H//16, H//32])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "338bbbba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2.48691e+00, 3.49555e+00, 5.12630e+00,  ..., 8.08980e-04, 7.84201e-03, 1.27847e-01],\n",
       "         [1.04869e+01, 3.49555e+00, 5.12630e+00,  ..., 8.08980e-04, 7.84201e-03, 1.27847e-01],\n",
       "         [1.84869e+01, 3.49555e+00, 5.12630e+00,  ..., 8.08980e-04, 7.84201e-03, 1.27847e-01],\n",
       "         ...,\n",
       "         [5.97028e+02, 3.53377e+02, 6.12005e+01,  ..., 1.01705e-02, 1.50065e-03, 3.06086e-03],\n",
       "         [6.15579e+02, 3.57104e+02, 5.20234e+01,  ..., 4.78162e-03, 2.56205e-03, 4.58687e-03],\n",
       "         [6.37331e+02, 3.67696e+02, 4.92063e+01,  ..., 6.70077e-03, 1.62391e-02, 9.34259e-03]]])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "e7360d87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2.48691e+00, 3.49555e+00, 5.12630e+00,  ..., 8.08982e-04, 7.84202e-03, 1.27847e-01],\n",
       "         [1.04869e+01, 3.49555e+00, 5.12630e+00,  ..., 8.08982e-04, 7.84202e-03, 1.27847e-01],\n",
       "         [1.84869e+01, 3.49555e+00, 5.12630e+00,  ..., 8.08982e-04, 7.84202e-03, 1.27847e-01],\n",
       "         ...,\n",
       "         [5.90872e+02, 3.74057e+02, 2.90891e+01,  ..., 2.20489e-03, 8.76354e-04, 8.86529e-04],\n",
       "         [6.23391e+02, 3.70629e+02, 3.82358e+01,  ..., 5.86631e-04, 4.67911e-04, 5.00718e-04],\n",
       "         [6.42960e+02, 3.75359e+02, 3.29330e+01,  ..., 2.74538e-03, 1.64567e-03, 2.05125e-03]]], device='cuda:0')"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee4f4be8",
   "metadata": {},
   "source": [
    "### Test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9b90ffac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.general import non_max_suppression, scale_coords, check_img_size\n",
    "from pathlib import Path\n",
    "from utils.loss import ComputeLoss\n",
    "from torch.cuda import amp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "172a5c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_model.eval().cuda()\n",
    "compute_loss = ComputeLoss(torch_model)  # init loss class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d8afbb8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "test:   0%|                                                                                                                           | 0/6467 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.60707], device='cuda:0')\n",
      "tensor([39784.87109], device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "save_hybrid = False\n",
    "cuda = 'cuda'\n",
    "device='cuda'\n",
    "conf_thres = 1e-4\n",
    "iou_thres = 0.6\n",
    "seen = 0\n",
    "scaler = amp.GradScaler(enabled=cuda)\n",
    "gs = max(int(torch_model.stride.max()), 32)  # grid size (max stride)\n",
    "imgsz, imgsz_test = [check_img_size(x, gs) for x in [640, 640]]  # verify imgsz are gs-multiples\n",
    "# test_path = data_dict['test']\n",
    "\n",
    "# train_loader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt=None,  # testloader\n",
    "#                                        hyp=hyp, cache=None, rect=True, rank=-1,\n",
    "#                                        world_size=1, workers=8,\n",
    "#                                        pad=0.5, prefix=colorstr('val: '))[0]\n",
    "\n",
    "for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc='test')):\n",
    "    nb, _, height, width = img.shape  # batch size, channels, height, width\n",
    "    img = img.to(device, non_blocking=True).float()/255.0\n",
    "    out, train_out = torch_model(img)  # forward\n",
    "    \n",
    "    targets[:, 2:] *= torch.Tensor([width, height, width, height]).to('cpu')  # to pixels\n",
    "    lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling\n",
    "    out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=True)\n",
    "    print(loss)\n",
    "    print(scaler.scale(loss))\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e07902b1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}