per_sample_gradients.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import TQS\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "EMBED_DIM = 64\n",
    "MAX_LENGTH = 100\n",
    "NUM_HEADS = 1\n",
    "NUM_LAYERS = 1\n",
    "DIM_FEEDFORWARD = 128\n",
    "TEST_LENGTH = 50\n",
    "TEST_BATCH = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_potentials = torch.randn(TEST_LENGTH, TEST_BATCH)\n",
    "test_spins = torch.zeros(TEST_LENGTH, TEST_BATCH)\n",
    "test_spins[torch.randint(0, TEST_LENGTH, (TEST_BATCH,)), torch.arange(TEST_BATCH)] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this situation each site has its own potential value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.8355, -1.7400,  1.0942,  ..., -1.1610,  1.1955, -1.3504],\n",
       "        [-0.7105, -0.3849,  1.4603,  ..., -0.0645,  0.0686, -0.0424],\n",
       "        [ 0.4510,  2.5088,  0.0482,  ...,  1.5008, -0.0663,  1.2951],\n",
       "        ...,\n",
       "        [-0.3750, -1.1446,  0.2271,  ..., -0.3643, -0.4005, -0.4082],\n",
       "        [ 0.1729, -0.7864,  0.2976,  ...,  1.2363,  2.0191,  0.8710],\n",
       "        [-0.5019,  0.4500, -0.5023,  ..., -0.2191,  0.0394,  0.5011]])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_potentials  # (seq_len, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 1.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.]])"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_spins  # (seq_len, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/spandan/anaconda3/envs/tqs2/lib/python3.12/site-packages/torch/nn/modules/transformer.py:306: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)\n",
      "  warnings.warn(f\"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}\")\n"
     ]
    }
   ],
   "source": [
    "tqs = TQS(\n",
    "    embed_dim=EMBED_DIM,\n",
    "    max_chain_len=MAX_LENGTH,\n",
    "    num_heads=NUM_HEADS,\n",
    "    num_layers=NUM_LAYERS,\n",
    "    possible_spins=2,\n",
    "    dim_feedforward=DIM_FEEDFORWARD,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "probs, phases = tqs(test_potentials, test_spins)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.4460, 0.4607],\n",
       "         [0.4032, 0.4227],\n",
       "         [0.5811, 0.5957],\n",
       "         ...,\n",
       "         [0.4324, 0.4341],\n",
       "         [0.6108, 0.5829],\n",
       "         [0.4041, 0.4245]],\n",
       "\n",
       "        [[0.4560, 0.4223],\n",
       "         [0.4975, 0.3897],\n",
       "         [0.5926, 0.5843],\n",
       "         ...,\n",
       "         [0.5123, 0.4368],\n",
       "         [0.5567, 0.5168],\n",
       "         [0.5307, 0.4366]],\n",
       "\n",
       "        [[0.5944, 0.5031],\n",
       "         [0.6034, 0.5769],\n",
       "         [0.5397, 0.4638],\n",
       "         ...,\n",
       "         [0.5894, 0.5567],\n",
       "         [0.5592, 0.4499],\n",
       "         [0.6072, 0.5547]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[0.6344, 0.4580],\n",
       "         [0.6461, 0.4298],\n",
       "         [0.6110, 0.4403],\n",
       "         ...,\n",
       "         [0.6812, 0.4494],\n",
       "         [0.6696, 0.4140],\n",
       "         [0.6364, 0.4233]],\n",
       "\n",
       "        [[0.6722, 0.4469],\n",
       "         [0.6540, 0.4431],\n",
       "         [0.6951, 0.4497],\n",
       "         ...,\n",
       "         [0.6424, 0.4508],\n",
       "         [0.6817, 0.4990],\n",
       "         [0.6902, 0.5106]],\n",
       "\n",
       "        [[0.7199, 0.4970],\n",
       "         [0.7130, 0.5040],\n",
       "         [0.7044, 0.5347],\n",
       "         ...,\n",
       "         [0.6890, 0.4737],\n",
       "         [0.7162, 0.5054],\n",
       "         [0.7174, 0.4782]]], grad_fn=<SigmoidBackward0>)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.5843],\n",
       "         [-0.8018],\n",
       "         [ 0.2631],\n",
       "         ...,\n",
       "         [-0.4593],\n",
       "         [ 0.2807],\n",
       "         [-0.8551]],\n",
       "\n",
       "        [[-0.2700],\n",
       "         [-0.0188],\n",
       "         [ 0.3330],\n",
       "         ...,\n",
       "         [ 0.4155],\n",
       "         [ 0.0548],\n",
       "         [ 0.2054]],\n",
       "\n",
       "        [[ 0.4977],\n",
       "         [ 0.5756],\n",
       "         [ 0.0881],\n",
       "         ...,\n",
       "         [ 0.4130],\n",
       "         [-0.6362],\n",
       "         [-0.2154]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 2.0218],\n",
       "         [ 2.1927],\n",
       "         [ 1.9539],\n",
       "         ...,\n",
       "         [ 2.1012],\n",
       "         [ 2.4222],\n",
       "         [ 2.2464]],\n",
       "\n",
       "        [[ 2.2077],\n",
       "         [ 2.1917],\n",
       "         [ 2.1699],\n",
       "         ...,\n",
       "         [ 2.0066],\n",
       "         [ 1.9795],\n",
       "         [ 1.8199]],\n",
       "\n",
       "        [[ 2.2033],\n",
       "         [ 1.9242],\n",
       "         [ 2.0259],\n",
       "         ...,\n",
       "         [ 2.2054],\n",
       "         [ 2.0765],\n",
       "         [ 1.8804]]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 1., 0.,  ..., 0., 0., 0.],\n",
       "        [1., 0., 0.,  ..., 1., 1., 1.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.]])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Spins sampled from the wave function the TQS models\n",
    "spins = tqs.sample_spins(test_potentials, tqs.max_chain_len)\n",
    "spins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 1.0\n",
    "psi_x, psi_l, psi_r = tqs.psi_terms(spins, test_potentials, T)\n",
    "E_loc = tqs.E_loc(psi_x, psi_l, psi_r, test_potentials, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([ 0.7700, -0.1845, -0.5371,  1.5401,  0.3018,  1.5673, -0.5706,  0.5321,\n",
      "        -0.0204, -1.1571,  1.4552,  1.0573,  0.3328,  1.2447,  0.1200,  0.9041,\n",
      "         0.1951, -0.8236,  0.8809,  0.8374, -0.3279,  0.5729, -0.2325,  0.4923,\n",
      "         1.0704,  0.6647, -0.4704, -0.8055, -1.3952,  0.6256,  0.2798,  1.0896,\n",
      "         0.4670,  0.0300, -0.9752,  3.0311,  2.4092,  0.3060, -0.0418, -1.2533,\n",
      "         0.5433, -1.1362,  1.2849,  0.8085, -0.4620,  0.3075, -0.4362,  1.1863,\n",
      "        -0.2932,  0.5470,  0.3605,  0.4797,  1.2236,  0.3412,  0.1848, -0.2307,\n",
      "         0.5630, -1.4219,  1.0084,  0.5428, -0.8511,  1.1395, -0.5453,  0.4003],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([-0.1294, -0.9766, -1.2813, -0.9127, -1.9270,  0.5791, -2.0543, -0.5088,\n",
      "        -1.1811,  0.9136, -2.1359, -0.4256, -0.2987, -0.1058,  0.7262,  0.4244,\n",
      "        -2.0061, -0.4241, -0.3564, -0.9538, -1.7064,  0.8925, -0.2464,  1.2790,\n",
      "         0.4191,  0.0980,  0.9135, -0.9046, -1.1397,  1.4422,  0.6763, -1.2613,\n",
      "        -1.5748,  0.2970,  1.3577, -0.0106, -0.0104,  1.3326, -1.5871, -0.0502,\n",
      "         1.0109, -0.8336, -0.5829,  0.0255, -0.6497, -1.0316,  1.5062,  0.6167,\n",
      "         0.2840, -0.6188,  1.1577,  0.3730,  1.3188, -1.4541, -0.7055,  0.0739,\n",
      "         0.6743, -1.1031,  0.5257, -0.1520, -1.4180, -0.6036,  1.0314,  0.5921],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-0.0325,  0.0317,  0.0013,  ...,  0.0096, -0.1311,  0.0082],\n",
      "        [-0.0445, -0.1377,  0.1477,  ...,  0.1227,  0.1165,  0.0879],\n",
      "        [ 0.0990,  0.1367,  0.1495,  ..., -0.0237, -0.1293,  0.1230],\n",
      "        ...,\n",
      "        [-0.0936, -0.0061,  0.0384,  ...,  0.1097, -0.0887, -0.0464],\n",
      "        [-0.0277, -0.0176,  0.1430,  ...,  0.1388, -0.1445,  0.0448],\n",
      "        [ 0.0570,  0.1527,  0.0655,  ..., -0.1247, -0.0883,  0.0490]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-1.1941e-01, -3.4386e-02, -1.6856e-02,  ..., -1.1792e-02,\n",
      "          8.2776e-02,  2.8318e-02],\n",
      "        [-8.2614e-02, -1.2132e-01,  1.0502e-01,  ...,  3.4139e-05,\n",
      "         -1.2312e-01,  3.6218e-02],\n",
      "        [-9.5822e-03,  6.0354e-02, -9.7457e-02,  ..., -3.9229e-02,\n",
      "          4.2696e-02,  6.5451e-02],\n",
      "        ...,\n",
      "        [-4.7530e-02,  5.4085e-02, -3.1076e-02,  ..., -8.6690e-02,\n",
      "         -4.5052e-02, -4.4182e-02],\n",
      "        [ 1.7151e-02,  9.3342e-02, -1.3774e-02,  ..., -8.6132e-02,\n",
      "         -1.5861e-02,  9.9138e-02],\n",
      "        [-5.5146e-02, -3.2271e-02, -7.9591e-02,  ..., -1.2299e-01,\n",
      "          1.1702e-01, -1.1056e-01]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.0778,  0.0391,  0.1179,  ...,  0.0116, -0.1213,  0.0859],\n",
      "        [ 0.0354, -0.0360, -0.0028,  ..., -0.0015,  0.0531,  0.0307],\n",
      "        [-0.0170, -0.0734,  0.0666,  ...,  0.1233, -0.0109, -0.0654],\n",
      "        ...,\n",
      "        [-0.0019,  0.0559, -0.0198,  ..., -0.1024, -0.0893, -0.0598],\n",
      "        [ 0.1128, -0.0201, -0.1078,  ...,  0.0272, -0.0865,  0.1053],\n",
      "        [-0.0458,  0.0905, -0.0589,  ..., -0.1032, -0.0435,  0.0810]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([ 7.5789e-02, -1.0873e-01,  1.1134e-01,  2.4613e-02,  1.2200e-01,\n",
      "         6.5407e-02, -1.8141e-02,  6.2868e-02,  7.5938e-02,  7.4019e-02,\n",
      "        -1.5137e-02, -1.2016e-01, -4.8110e-03, -1.0832e-01,  1.1633e-01,\n",
      "        -3.5827e-02, -7.4036e-03, -8.8470e-02,  1.0862e-01,  6.2222e-02,\n",
      "         6.5405e-02, -5.9155e-02, -8.1863e-02,  7.5206e-05,  4.4275e-02,\n",
      "         1.1219e-01, -1.1101e-01, -1.2356e-02,  3.3764e-02,  4.7628e-02,\n",
      "         9.4678e-02, -5.8971e-02, -4.5981e-02, -9.1465e-02,  6.8468e-02,\n",
      "         3.8986e-02, -8.6881e-02,  3.3256e-02,  5.8579e-02,  1.1930e-01,\n",
      "        -1.2064e-01,  5.3566e-02,  2.7804e-02,  3.0437e-02, -4.9965e-02,\n",
      "        -1.2483e-01,  1.0828e-01,  4.1472e-02,  8.0523e-02,  8.3026e-02,\n",
      "         1.1542e-01, -4.6775e-02,  5.3472e-02,  6.6470e-03,  6.1962e-02,\n",
      "        -5.6780e-02, -1.1705e-01, -3.8531e-03, -2.2056e-02, -5.3301e-02,\n",
      "         3.4090e-02, -6.9624e-02,  2.1264e-02,  3.1408e-02, -5.7265e-02,\n",
      "        -7.0230e-02, -6.7608e-03,  1.0030e-01,  4.7230e-02, -2.4996e-02,\n",
      "        -1.1278e-01, -9.8338e-02,  5.2228e-03,  7.3433e-02, -1.1254e-01,\n",
      "        -5.4799e-02,  9.7720e-02, -8.3697e-02, -6.1971e-02,  6.0909e-02,\n",
      "         8.8443e-02,  1.1619e-01,  1.1551e-02, -1.1984e-01, -4.2313e-02,\n",
      "         3.4900e-02,  1.0260e-01,  1.0316e-01, -1.1678e-02, -6.9727e-02,\n",
      "        -7.7097e-02,  1.1701e-02, -1.3508e-02, -1.2013e-01, -4.7125e-02,\n",
      "         4.6152e-02, -1.0063e-01, -1.0096e-01,  8.4908e-03,  3.0550e-02,\n",
      "        -6.0123e-02,  1.2334e-01, -1.2115e-01, -1.0311e-01,  1.7173e-03,\n",
      "         2.4510e-02,  3.7142e-02, -3.9464e-02, -1.0750e-01, -9.1214e-02,\n",
      "        -4.6268e-02,  1.0808e-01,  6.0863e-02, -1.9604e-02,  1.9614e-02,\n",
      "         1.7090e-02,  3.3229e-02, -2.3045e-02,  9.3213e-02, -1.1551e-01,\n",
      "         1.8696e-02, -1.2199e-01,  2.2935e-02,  1.1038e-01, -1.0004e-01,\n",
      "         8.4598e-02, -1.2622e-02, -9.7852e-03], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.0050,  0.0240, -0.0076,  ..., -0.0042,  0.0876,  0.0790],\n",
      "        [ 0.0287, -0.0250, -0.0312,  ...,  0.0361, -0.0525, -0.0806],\n",
      "        [-0.0348,  0.0750,  0.0333,  ..., -0.0716,  0.0185,  0.0322],\n",
      "        ...,\n",
      "        [ 0.0253, -0.0633,  0.0679,  ...,  0.0161, -0.0536,  0.0142],\n",
      "        [-0.0398, -0.0214,  0.0680,  ...,  0.0417, -0.0672, -0.0720],\n",
      "        [-0.0587,  0.0088, -0.0351,  ..., -0.0780,  0.0242,  0.0265]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([-0.0676, -0.0312, -0.0687, -0.0385, -0.0756, -0.0576, -0.0531,  0.0734,\n",
      "         0.0285,  0.0575, -0.0281, -0.0524, -0.0660, -0.0458,  0.0839, -0.0127,\n",
      "        -0.0417,  0.0325,  0.0046, -0.0326, -0.0261,  0.0090, -0.0417, -0.0149,\n",
      "        -0.0032, -0.0179, -0.0167,  0.0577,  0.0812,  0.0534, -0.0444,  0.0661,\n",
      "         0.0786, -0.0275, -0.0694, -0.0127,  0.0107, -0.0841, -0.0072,  0.0545,\n",
      "        -0.0185,  0.0309,  0.0132,  0.0382,  0.0481,  0.0650, -0.0101,  0.0378,\n",
      "         0.0768, -0.0723, -0.0586,  0.0834,  0.0428,  0.0804, -0.0769,  0.0824,\n",
      "         0.0073,  0.0407,  0.0258,  0.0804,  0.0596,  0.0513, -0.0064,  0.0262],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-0.0325,  0.0317,  0.0013,  ...,  0.0096, -0.1311,  0.0082],\n",
      "        [-0.0445, -0.1377,  0.1477,  ...,  0.1227,  0.1165,  0.0879],\n",
      "        [ 0.0990,  0.1367,  0.1495,  ..., -0.0237, -0.1293,  0.1230],\n",
      "        ...,\n",
      "        [-0.0936, -0.0061,  0.0384,  ...,  0.1097, -0.0887, -0.0464],\n",
      "        [-0.0277, -0.0176,  0.1430,  ...,  0.1388, -0.1445,  0.0448],\n",
      "        [ 0.0570,  0.1527,  0.0655,  ..., -0.1247, -0.0883,  0.0490]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-1.1941e-01, -3.4386e-02, -1.6856e-02,  ..., -1.1792e-02,\n",
      "          8.2776e-02,  2.8318e-02],\n",
      "        [-8.2614e-02, -1.2132e-01,  1.0502e-01,  ...,  3.4139e-05,\n",
      "         -1.2312e-01,  3.6218e-02],\n",
      "        [-9.5822e-03,  6.0354e-02, -9.7457e-02,  ..., -3.9229e-02,\n",
      "          4.2696e-02,  6.5451e-02],\n",
      "        ...,\n",
      "        [-4.7530e-02,  5.4085e-02, -3.1076e-02,  ..., -8.6690e-02,\n",
      "         -4.5052e-02, -4.4182e-02],\n",
      "        [ 1.7151e-02,  9.3342e-02, -1.3774e-02,  ..., -8.6132e-02,\n",
      "         -1.5861e-02,  9.9138e-02],\n",
      "        [-5.5146e-02, -3.2271e-02, -7.9591e-02,  ..., -1.2299e-01,\n",
      "          1.1702e-01, -1.1056e-01]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.0778,  0.0391,  0.1179,  ...,  0.0116, -0.1213,  0.0859],\n",
      "        [ 0.0354, -0.0360, -0.0028,  ..., -0.0015,  0.0531,  0.0307],\n",
      "        [-0.0170, -0.0734,  0.0666,  ...,  0.1233, -0.0109, -0.0654],\n",
      "        ...,\n",
      "        [-0.0019,  0.0559, -0.0198,  ..., -0.1024, -0.0893, -0.0598],\n",
      "        [ 0.1128, -0.0201, -0.1078,  ...,  0.0272, -0.0865,  0.1053],\n",
      "        [-0.0458,  0.0905, -0.0589,  ..., -0.1032, -0.0435,  0.0810]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([ 7.5789e-02, -1.0873e-01,  1.1134e-01,  2.4613e-02,  1.2200e-01,\n",
      "         6.5407e-02, -1.8141e-02,  6.2868e-02,  7.5938e-02,  7.4019e-02,\n",
      "        -1.5137e-02, -1.2016e-01, -4.8110e-03, -1.0832e-01,  1.1633e-01,\n",
      "        -3.5827e-02, -7.4036e-03, -8.8470e-02,  1.0862e-01,  6.2222e-02,\n",
      "         6.5405e-02, -5.9155e-02, -8.1863e-02,  7.5206e-05,  4.4275e-02,\n",
      "         1.1219e-01, -1.1101e-01, -1.2356e-02,  3.3764e-02,  4.7628e-02,\n",
      "         9.4678e-02, -5.8971e-02, -4.5981e-02, -9.1465e-02,  6.8468e-02,\n",
      "         3.8986e-02, -8.6881e-02,  3.3256e-02,  5.8579e-02,  1.1930e-01,\n",
      "        -1.2064e-01,  5.3566e-02,  2.7804e-02,  3.0437e-02, -4.9965e-02,\n",
      "        -1.2483e-01,  1.0828e-01,  4.1472e-02,  8.0523e-02,  8.3026e-02,\n",
      "         1.1542e-01, -4.6775e-02,  5.3472e-02,  6.6470e-03,  6.1962e-02,\n",
      "        -5.6780e-02, -1.1705e-01, -3.8531e-03, -2.2056e-02, -5.3301e-02,\n",
      "         3.4090e-02, -6.9624e-02,  2.1264e-02,  3.1408e-02, -5.7265e-02,\n",
      "        -7.0230e-02, -6.7608e-03,  1.0030e-01,  4.7230e-02, -2.4996e-02,\n",
      "        -1.1278e-01, -9.8338e-02,  5.2228e-03,  7.3433e-02, -1.1254e-01,\n",
      "        -5.4799e-02,  9.7720e-02, -8.3697e-02, -6.1971e-02,  6.0909e-02,\n",
      "         8.8443e-02,  1.1619e-01,  1.1551e-02, -1.1984e-01, -4.2313e-02,\n",
      "         3.4900e-02,  1.0260e-01,  1.0316e-01, -1.1678e-02, -6.9727e-02,\n",
      "        -7.7097e-02,  1.1701e-02, -1.3508e-02, -1.2013e-01, -4.7125e-02,\n",
      "         4.6152e-02, -1.0063e-01, -1.0096e-01,  8.4908e-03,  3.0550e-02,\n",
      "        -6.0123e-02,  1.2334e-01, -1.2115e-01, -1.0311e-01,  1.7173e-03,\n",
      "         2.4510e-02,  3.7142e-02, -3.9464e-02, -1.0750e-01, -9.1214e-02,\n",
      "        -4.6268e-02,  1.0808e-01,  6.0863e-02, -1.9604e-02,  1.9614e-02,\n",
      "         1.7090e-02,  3.3229e-02, -2.3045e-02,  9.3213e-02, -1.1551e-01,\n",
      "         1.8696e-02, -1.2199e-01,  2.2935e-02,  1.1038e-01, -1.0004e-01,\n",
      "         8.4598e-02, -1.2622e-02, -9.7852e-03], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.0050,  0.0240, -0.0076,  ..., -0.0042,  0.0876,  0.0790],\n",
      "        [ 0.0287, -0.0250, -0.0312,  ...,  0.0361, -0.0525, -0.0806],\n",
      "        [-0.0348,  0.0750,  0.0333,  ..., -0.0716,  0.0185,  0.0322],\n",
      "        ...,\n",
      "        [ 0.0253, -0.0633,  0.0679,  ...,  0.0161, -0.0536,  0.0142],\n",
      "        [-0.0398, -0.0214,  0.0680,  ...,  0.0417, -0.0672, -0.0720],\n",
      "        [-0.0587,  0.0088, -0.0351,  ..., -0.0780,  0.0242,  0.0265]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([-0.0676, -0.0312, -0.0687, -0.0385, -0.0756, -0.0576, -0.0531,  0.0734,\n",
      "         0.0285,  0.0575, -0.0281, -0.0524, -0.0660, -0.0458,  0.0839, -0.0127,\n",
      "        -0.0417,  0.0325,  0.0046, -0.0326, -0.0261,  0.0090, -0.0417, -0.0149,\n",
      "        -0.0032, -0.0179, -0.0167,  0.0577,  0.0812,  0.0534, -0.0444,  0.0661,\n",
      "         0.0786, -0.0275, -0.0694, -0.0127,  0.0107, -0.0841, -0.0072,  0.0545,\n",
      "        -0.0185,  0.0309,  0.0132,  0.0382,  0.0481,  0.0650, -0.0101,  0.0378,\n",
      "         0.0768, -0.0723, -0.0586,  0.0834,  0.0428,  0.0804, -0.0769,  0.0824,\n",
      "         0.0073,  0.0407,  0.0258,  0.0804,  0.0596,  0.0513, -0.0064,  0.0262],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,\n",
      "        1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.0690,  0.0121, -0.0915,  0.0810,  0.0350, -0.0539,  0.0978, -0.0830,\n",
      "         -0.0174, -0.0597, -0.1182,  0.0316,  0.0819, -0.0258,  0.1205, -0.0850,\n",
      "         -0.0551,  0.0983, -0.0624, -0.0090, -0.0534,  0.0628, -0.0991, -0.0277,\n",
      "         -0.0306, -0.0585,  0.0698,  0.0176,  0.0395, -0.0027,  0.1223,  0.1039,\n",
      "         -0.0618,  0.0936, -0.0623,  0.0283, -0.0016, -0.0571,  0.0946,  0.0964,\n",
      "         -0.0558, -0.1117,  0.0825,  0.0259, -0.0432,  0.0726, -0.0736,  0.0291,\n",
      "          0.0587,  0.0887, -0.0390,  0.0329,  0.0451, -0.0013, -0.0275, -0.1062,\n",
      "         -0.0493, -0.0969, -0.1074, -0.1143, -0.0176,  0.0734,  0.0896,  0.0661],\n",
      "        [-0.0367,  0.0962, -0.0982, -0.0030, -0.0886, -0.0173,  0.1184, -0.0084,\n",
      "          0.1029, -0.0254, -0.0594,  0.0843, -0.0629, -0.0993, -0.0114, -0.0339,\n",
      "         -0.0252,  0.0487, -0.0367,  0.0404,  0.0817,  0.0974, -0.0860, -0.0286,\n",
      "          0.0192,  0.0502, -0.1021, -0.0415, -0.0581, -0.0935, -0.1093,  0.0923,\n",
      "         -0.0869, -0.0074, -0.0612,  0.0156,  0.0858, -0.0422, -0.0091,  0.1025,\n",
      "          0.0484, -0.0553,  0.0183,  0.0666,  0.1031, -0.0862, -0.0277,  0.0680,\n",
      "          0.0415, -0.0668,  0.0005,  0.0346,  0.0410, -0.1201, -0.0105,  0.1124,\n",
      "         -0.0631, -0.0479,  0.0219, -0.0469,  0.1062,  0.0313,  0.0709,  0.0228]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0.0308, 0.0695], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-0.0514,  0.0039,  0.0413,  0.0790,  0.0715,  0.1192, -0.0756, -0.0523,\n",
      "          0.1233, -0.0732,  0.0866,  0.0336,  0.0482, -0.0267, -0.0848, -0.0850,\n",
      "         -0.0587, -0.1232, -0.0661, -0.0562, -0.0278, -0.1070,  0.1220, -0.0993,\n",
      "          0.0713,  0.0399,  0.0222,  0.0820,  0.0669, -0.1018,  0.0885, -0.0718,\n",
      "          0.0713,  0.0404, -0.1113, -0.0180, -0.0455, -0.0247, -0.0627, -0.0774,\n",
      "         -0.1007, -0.0524, -0.1171, -0.0096,  0.0817,  0.0609, -0.1241, -0.1037,\n",
      "          0.0011, -0.0931,  0.0969, -0.0094, -0.0850,  0.0948, -0.0738, -0.1157,\n",
      "         -0.1201,  0.0584,  0.0157, -0.0768, -0.1065,  0.0339, -0.0961,  0.0664]],\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([0.0783], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "for i in tqs.parameters():\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Per-Sample Gradients\n",
    "\n",
    "The target is\n",
    "\n",
    "$$\\frac{d}{d\\Theta}\\left ( P(x; \\theta)\\right )$$\n",
    "\n",
    "\n",
    "...which ends up being represented as some tuple of gradients of $P$ with respect to $\\theta_1, \\theta_2, ...$:\n",
    "\n",
    "$$\\left ( \\frac{dP}{d\\theta_1}, \\frac{dP}{d\\theta_2}, ..., \\frac{dP}{d\\theta_p} \\right )$$\n",
    "\n",
    "...assuming we have $p$ model parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Simple CNN Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From https://pytorch.org/tutorials/intermediate/per_sample_grads.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output\n",
    "\n",
    "\n",
    "def loss_fn(predictions, targets):\n",
    "    return F.nll_loss(predictions, targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\"\n",
    "\n",
    "num_models = 10\n",
    "batch_size = 64\n",
    "data = torch.randn(batch_size, 1, 28, 28, device=device)\n",
    "\n",
    "targets = torch.randint(10, (64,), device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3039, device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Simple training computes an average across a batch\n",
    "model = SimpleCNN().to(device=device)\n",
    "predictions = model(data)  # move the entire mini-batch through the model\n",
    "\n",
    "loss = loss_fn(predictions, targets)\n",
    "loss  # Averaged over the batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()  # back propagate the 'average' gradient of this mini-batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of computing an averaged loss, we'd like to separate the loss function's computation graph (and thus its loss gradient, post backprop) between each sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.func import functional_call, vmap, grad\n",
    "\n",
    "params = {k: v.detach() for k, v in model.named_parameters()}\n",
    "buffers = {k: v.detach() for k, v in model.named_buffers()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss(params, buffers, sample, target):\n",
    "    \"\"\"\n",
    "    Computes loss gradients for a single sample with respect to\n",
    "    the model parameters (affected by differentiation) and buffers\n",
    "    (not affected by differentiation).\n",
    "    \"\"\"\n",
    "    batch = sample.unsqueeze(0)\n",
    "    targets = target.unsqueeze(0)\n",
    "\n",
    "    # Swaps model parameters and buffers with the ones provided,\n",
    "    # using a name -> tensor (or value) mapping\n",
    "    predictions = functional_call(model, (params, buffers), (batch,))\n",
    "    loss = loss_fn(predictions, targets)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arguments of `functional_call`:\n",
    "1. The module to transplant the parameters and buffers into\n",
    "2. The parameters and buffers themselves (named, using a dictionary, so that PyTorch knows where each one goes)\n",
    "3. The arguments to the function--in this case, the entire \"batch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[[[-0.1084,  0.0424, -0.3139],\n",
      "          [ 0.1655, -0.1783,  0.2891],\n",
      "          [ 0.2499,  0.0475,  0.1136]]],\n",
      "\n",
      "\n",
      "        [[[-0.1207, -0.1051, -0.1374],\n",
      "          [-0.0912,  0.1755, -0.0016],\n",
      "          [-0.0721,  0.1740, -0.0377]]],\n",
      "\n",
      "\n",
      "        [[[-0.0946, -0.2715,  0.3007],\n",
      "          [ 0.2387,  0.0555, -0.1468],\n",
      "          [-0.2477,  0.2970,  0.1652]]],\n",
      "\n",
      "\n",
      "        [[[-0.0218,  0.0265, -0.0242],\n",
      "          [ 0.0138, -0.2554,  0.2154],\n",
      "          [-0.0746,  0.0563,  0.1636]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1062,  0.2369,  0.1340],\n",
      "          [-0.1994,  0.0486, -0.0190],\n",
      "          [ 0.2310,  0.1852, -0.0057]]],\n",
      "\n",
      "\n",
      "        [[[-0.2098, -0.1552,  0.0312],\n",
      "          [ 0.2002,  0.2086, -0.0492],\n",
      "          [-0.1418,  0.1076, -0.1305]]],\n",
      "\n",
      "\n",
      "        [[[ 0.2461,  0.2975, -0.0725],\n",
      "          [-0.2895,  0.1223, -0.1883],\n",
      "          [-0.1135, -0.1375,  0.2733]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1612,  0.2042, -0.1575],\n",
      "          [ 0.2084, -0.3041, -0.0166],\n",
      "          [-0.1108,  0.1487,  0.0416]]],\n",
      "\n",
      "\n",
      "        [[[-0.0860,  0.1731,  0.0781],\n",
      "          [-0.1788,  0.1727,  0.1643],\n",
      "          [ 0.2228,  0.1225,  0.3109]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0876,  0.0007,  0.0502],\n",
      "          [ 0.1651,  0.1885,  0.1416],\n",
      "          [ 0.1714,  0.0746, -0.1605]]],\n",
      "\n",
      "\n",
      "        [[[-0.1184, -0.0699,  0.0863],\n",
      "          [-0.2824,  0.0130, -0.3175],\n",
      "          [ 0.3190, -0.1806,  0.2634]]],\n",
      "\n",
      "\n",
      "        [[[-0.1832,  0.2853, -0.0408],\n",
      "          [-0.3036, -0.0914, -0.1076],\n",
      "          [-0.2646,  0.1378,  0.2598]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1710, -0.0300, -0.0533],\n",
      "          [-0.0528,  0.2133, -0.0243],\n",
      "          [-0.3220,  0.0132, -0.0641]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0099,  0.2895,  0.0917],\n",
      "          [ 0.1722, -0.1846, -0.0589],\n",
      "          [-0.2416, -0.1231,  0.2799]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0091, -0.1219,  0.2342],\n",
      "          [ 0.0209,  0.2767,  0.1340],\n",
      "          [-0.3186,  0.1719,  0.3191]]],\n",
      "\n",
      "\n",
      "        [[[-0.1264, -0.0947,  0.0806],\n",
      "          [ 0.0421,  0.2946, -0.2986],\n",
      "          [-0.2788,  0.2291,  0.1660]]],\n",
      "\n",
      "\n",
      "        [[[-0.2786, -0.0621, -0.3195],\n",
      "          [-0.0674, -0.0630,  0.0060],\n",
      "          [-0.2711, -0.0371,  0.0818]]],\n",
      "\n",
      "\n",
      "        [[[-0.2440,  0.0168, -0.2356],\n",
      "          [ 0.2044, -0.1151, -0.1611],\n",
      "          [-0.2293,  0.1449, -0.2304]]],\n",
      "\n",
      "\n",
      "        [[[ 0.2063,  0.2947, -0.3299],\n",
      "          [-0.1287,  0.0851, -0.0837],\n",
      "          [-0.0686,  0.1300,  0.2938]]],\n",
      "\n",
      "\n",
      "        [[[-0.3096,  0.1436,  0.1071],\n",
      "          [-0.1070, -0.1141,  0.0476],\n",
      "          [-0.1887,  0.0285, -0.2690]]],\n",
      "\n",
      "\n",
      "        [[[-0.0637,  0.2782, -0.3018],\n",
      "          [-0.0498, -0.1136,  0.1018],\n",
      "          [-0.2118,  0.0768, -0.2642]]],\n",
      "\n",
      "\n",
      "        [[[-0.1788,  0.2215, -0.1503],\n",
      "          [ 0.1783, -0.1960,  0.2056],\n",
      "          [ 0.3205,  0.2245, -0.0220]]],\n",
      "\n",
      "\n",
      "        [[[-0.1503, -0.1933,  0.3251],\n",
      "          [ 0.0342,  0.2867, -0.3114],\n",
      "          [ 0.1686, -0.2938,  0.1935]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0764,  0.1853,  0.2425],\n",
      "          [ 0.0801,  0.0348, -0.2382],\n",
      "          [ 0.2219, -0.0374, -0.3265]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1366, -0.3115,  0.1429],\n",
      "          [ 0.0910,  0.0487, -0.0660],\n",
      "          [ 0.1346, -0.2557,  0.3127]]],\n",
      "\n",
      "\n",
      "        [[[-0.1727,  0.2936, -0.3200],\n",
      "          [-0.2767, -0.1815, -0.0955],\n",
      "          [-0.0967, -0.2507, -0.2792]]],\n",
      "\n",
      "\n",
      "        [[[ 0.2503, -0.1005, -0.0823],\n",
      "          [-0.1077,  0.1695,  0.2918],\n",
      "          [ 0.1806,  0.2922, -0.0921]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0003, -0.0028,  0.0451],\n",
      "          [ 0.0254, -0.0841,  0.1934],\n",
      "          [ 0.0119, -0.0586,  0.3304]]],\n",
      "\n",
      "\n",
      "        [[[-0.3108,  0.2535, -0.2508],\n",
      "          [ 0.0478,  0.2614,  0.1835],\n",
      "          [-0.2976,  0.1228,  0.0274]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1999, -0.0818,  0.1843],\n",
      "          [-0.1047, -0.1606,  0.1539],\n",
      "          [ 0.0816,  0.0009,  0.2956]]],\n",
      "\n",
      "\n",
      "        [[[-0.2338, -0.1855, -0.2726],\n",
      "          [ 0.0191,  0.2090, -0.2895],\n",
      "          [ 0.2704, -0.3068,  0.0519]]],\n",
      "\n",
      "\n",
      "        [[[-0.2445, -0.1509, -0.1576],\n",
      "          [ 0.1453,  0.0736, -0.0322],\n",
      "          [-0.0682,  0.1615, -0.0276]]]], device='cuda:0', requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([ 0.1176,  0.2791,  0.0127, -0.2463, -0.0178, -0.1985, -0.0314,  0.1866,\n",
      "        -0.1807, -0.1390, -0.2774, -0.1075, -0.1732,  0.2612, -0.1149, -0.0391,\n",
      "        -0.2490,  0.0519, -0.2534,  0.0720,  0.1817, -0.2042, -0.3159, -0.0778,\n",
      "         0.1779, -0.2923, -0.2999, -0.1905,  0.1117, -0.0645,  0.0614, -0.0592],\n",
      "       device='cuda:0', requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[[[-2.4306e-02, -1.8751e-02, -2.0324e-03],\n",
      "          [ 1.3566e-02,  4.8158e-02,  4.9633e-02],\n",
      "          [-1.4187e-03, -5.0378e-02,  5.0223e-03]],\n",
      "\n",
      "         [[ 3.9221e-02,  3.2672e-02, -4.0997e-02],\n",
      "          [-2.4202e-02, -2.4069e-02,  1.9433e-02],\n",
      "          [-5.3987e-02, -2.2816e-02,  2.6622e-03]],\n",
      "\n",
      "         [[ 5.4145e-02,  3.6966e-02, -5.5133e-02],\n",
      "          [-3.5722e-02, -4.1347e-02, -4.3220e-02],\n",
      "          [ 4.2077e-02, -1.6663e-02, -4.5847e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-5.7426e-02, -5.2370e-02,  9.6920e-03],\n",
      "          [-5.4665e-02, -5.5117e-02, -5.8737e-02],\n",
      "          [ 2.2378e-03, -2.5397e-02,  5.7548e-02]],\n",
      "\n",
      "         [[ 4.5802e-02,  9.8386e-03, -2.8811e-02],\n",
      "          [-5.8856e-02, -5.1607e-02,  2.2990e-02],\n",
      "          [ 4.4320e-02,  1.1785e-02, -2.0781e-02]],\n",
      "\n",
      "         [[ 2.0348e-03, -4.5806e-02,  1.5470e-02],\n",
      "          [ 1.9497e-02,  3.4176e-02,  3.6330e-02],\n",
      "          [-3.7492e-02,  4.8235e-02, -4.1963e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 2.4407e-02,  7.5738e-04, -4.4577e-03],\n",
      "          [-3.8982e-02,  5.7839e-02, -7.0807e-03],\n",
      "          [-2.2986e-02, -2.6743e-02,  1.3247e-02]],\n",
      "\n",
      "         [[ 3.0260e-02, -5.5304e-03,  4.0384e-02],\n",
      "          [-3.5988e-03, -7.6780e-03,  8.6703e-03],\n",
      "          [ 1.0672e-02, -5.1054e-02, -1.6626e-02]],\n",
      "\n",
      "         [[ 1.1365e-02, -2.4209e-02, -1.9724e-03],\n",
      "          [ 1.4473e-02, -2.3430e-02,  4.9205e-02],\n",
      "          [ 1.1180e-02,  2.3083e-02,  1.0016e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-2.6666e-03, -2.2539e-02, -2.2991e-02],\n",
      "          [ 6.4390e-03,  3.3677e-03,  3.7940e-02],\n",
      "          [ 6.9841e-03, -3.6802e-02,  5.8537e-02]],\n",
      "\n",
      "         [[ 2.9017e-02,  4.6025e-02, -6.4235e-03],\n",
      "          [-3.4744e-02, -9.1340e-03, -5.0396e-02],\n",
      "          [-3.8007e-03, -4.2183e-02,  4.3246e-02]],\n",
      "\n",
      "         [[ 1.6527e-03,  3.0006e-02,  5.7692e-02],\n",
      "          [-2.4401e-02, -1.8890e-02, -5.4968e-03],\n",
      "          [ 2.1378e-02, -4.4616e-02, -4.1644e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 4.0031e-02,  5.8591e-02, -8.6444e-03],\n",
      "          [-4.4695e-03, -5.3456e-02,  1.0663e-02],\n",
      "          [ 4.9043e-02,  3.6082e-03,  1.7605e-03]],\n",
      "\n",
      "         [[-5.5848e-02, -1.4519e-02,  9.0456e-03],\n",
      "          [ 2.3658e-02,  4.6769e-02,  4.9199e-02],\n",
      "          [-3.1785e-02, -7.4247e-03, -2.6327e-03]],\n",
      "\n",
      "         [[ 3.8800e-02,  3.5541e-02, -5.4504e-02],\n",
      "          [-1.1235e-02,  3.5460e-02, -3.3841e-03],\n",
      "          [ 8.5156e-03,  4.9525e-02,  9.1458e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-5.8589e-02, -7.9952e-03, -4.2556e-02],\n",
      "          [ 4.7465e-02, -5.6107e-02, -5.5800e-02],\n",
      "          [ 3.3444e-02,  5.6570e-02,  3.3498e-02]],\n",
      "\n",
      "         [[-1.1178e-02, -3.1207e-02, -9.0059e-03],\n",
      "          [-8.4177e-03, -6.1694e-03,  2.5947e-03],\n",
      "          [-7.7374e-03,  2.7359e-02,  2.8699e-02]],\n",
      "\n",
      "         [[ 3.3574e-02, -4.5986e-03, -3.6012e-02],\n",
      "          [-2.4177e-02, -7.5374e-03, -4.3222e-03],\n",
      "          [ 1.5222e-02,  5.0746e-03,  1.4577e-02]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 5.0525e-02, -1.3709e-02,  5.8846e-02],\n",
      "          [ 4.8263e-02,  5.4800e-02, -1.9747e-02],\n",
      "          [-4.1727e-02,  3.2882e-02,  9.3166e-03]],\n",
      "\n",
      "         [[ 2.4836e-02,  1.1619e-02,  1.1596e-02],\n",
      "          [ 6.4268e-03, -3.0227e-02,  1.4438e-02],\n",
      "          [-6.9298e-03, -4.9661e-02, -1.9881e-02]],\n",
      "\n",
      "         [[-5.3959e-02,  4.4565e-02,  3.8684e-03],\n",
      "          [-1.7676e-02,  2.5185e-02, -5.7895e-02],\n",
      "          [ 2.5201e-02,  2.8994e-02, -1.2344e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 4.4113e-02, -4.6908e-02,  3.7197e-02],\n",
      "          [ 1.8132e-02, -1.2989e-02, -4.8867e-02],\n",
      "          [ 6.1487e-03,  1.9476e-02, -1.5290e-03]],\n",
      "\n",
      "         [[-5.5071e-02, -3.0257e-02, -4.0586e-02],\n",
      "          [-8.6750e-03,  3.1739e-02, -1.1758e-02],\n",
      "          [-2.9693e-03,  3.2453e-02, -2.5182e-02]],\n",
      "\n",
      "         [[-4.2898e-02, -5.3553e-02,  5.3072e-02],\n",
      "          [ 8.4294e-03,  1.5923e-02,  4.7231e-02],\n",
      "          [-4.8105e-03, -9.5515e-04, -3.8464e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 9.2453e-03, -3.2662e-02, -5.7574e-02],\n",
      "          [ 2.1783e-02,  1.4709e-02, -2.3050e-02],\n",
      "          [ 5.8672e-02,  4.6333e-02, -9.4220e-03]],\n",
      "\n",
      "         [[-7.0153e-05,  2.1597e-02, -2.1167e-02],\n",
      "          [-1.7748e-02, -3.2309e-02,  1.6283e-02],\n",
      "          [ 2.8621e-03,  2.0096e-03,  5.0500e-02]],\n",
      "\n",
      "         [[-9.6830e-03,  1.1882e-02, -4.3738e-02],\n",
      "          [-4.2280e-03, -1.4669e-04, -3.6115e-02],\n",
      "          [ 2.9424e-02, -2.4835e-02,  3.5561e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-4.0143e-02,  2.4223e-02,  2.7427e-02],\n",
      "          [-2.1188e-02,  6.7718e-03,  3.4459e-02],\n",
      "          [-2.9381e-02,  6.3398e-03, -9.3170e-04]],\n",
      "\n",
      "         [[ 5.3007e-02, -6.8124e-04,  4.6470e-02],\n",
      "          [ 1.8348e-02, -4.9496e-02, -3.8624e-02],\n",
      "          [-5.6020e-02, -4.4637e-02,  1.2103e-03]],\n",
      "\n",
      "         [[-6.1402e-03,  4.8870e-02,  5.3060e-02],\n",
      "          [-2.3468e-02, -2.9031e-02, -3.1038e-02],\n",
      "          [-6.7096e-03, -4.4595e-02,  4.2470e-02]]],\n",
      "\n",
      "\n",
      "        [[[-2.4492e-02, -5.2172e-02,  5.6049e-02],\n",
      "          [ 3.3522e-02,  5.4762e-03,  3.5252e-02],\n",
      "          [ 3.4880e-03,  2.6319e-02, -1.3770e-02]],\n",
      "\n",
      "         [[ 4.0110e-03,  2.1153e-02,  4.9358e-02],\n",
      "          [ 5.1858e-02,  2.7305e-02,  2.7474e-02],\n",
      "          [-5.8289e-02,  2.9184e-02,  1.6084e-02]],\n",
      "\n",
      "         [[ 2.0692e-02,  3.8651e-02, -5.8524e-02],\n",
      "          [-5.4035e-02,  1.1142e-02,  5.4433e-04],\n",
      "          [ 3.5670e-02,  2.3305e-02,  3.9735e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-4.0456e-04, -2.4339e-02, -3.1984e-02],\n",
      "          [-4.0253e-03,  3.6734e-02, -5.5210e-02],\n",
      "          [-2.4478e-02, -4.8022e-02, -6.5821e-03]],\n",
      "\n",
      "         [[ 1.6952e-02,  5.5572e-02,  4.8609e-02],\n",
      "          [-2.0176e-03, -4.9891e-02, -3.7121e-02],\n",
      "          [-4.5094e-02, -2.1985e-02, -3.3074e-02]],\n",
      "\n",
      "         [[-5.4968e-02, -4.7646e-02, -4.6278e-02],\n",
      "          [ 3.4342e-02, -2.9623e-02,  3.1922e-02],\n",
      "          [-1.5788e-02,  4.4650e-02,  4.8803e-02]]]], device='cuda:0',\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([-0.0137,  0.0240,  0.0152,  0.0269, -0.0578, -0.0585, -0.0030,  0.0273,\n",
      "        -0.0423, -0.0511,  0.0585, -0.0284, -0.0411,  0.0270, -0.0139, -0.0366,\n",
      "         0.0429, -0.0149,  0.0513,  0.0224, -0.0382,  0.0075,  0.0317, -0.0515,\n",
      "        -0.0315, -0.0463,  0.0219,  0.0485, -0.0471,  0.0322, -0.0178,  0.0346,\n",
      "         0.0341,  0.0023,  0.0053, -0.0223,  0.0116, -0.0016,  0.0056,  0.0030,\n",
      "         0.0024,  0.0117,  0.0234,  0.0393, -0.0144, -0.0353,  0.0018, -0.0387,\n",
      "        -0.0132,  0.0013, -0.0286,  0.0418, -0.0557,  0.0107, -0.0440, -0.0406,\n",
      "         0.0556, -0.0554, -0.0112, -0.0480, -0.0117,  0.0573,  0.0217,  0.0237],\n",
      "       device='cuda:0', requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.0070,  0.0043,  0.0024,  ..., -0.0030, -0.0046,  0.0025],\n",
      "        [ 0.0013, -0.0029,  0.0073,  ...,  0.0044, -0.0101,  0.0078],\n",
      "        [ 0.0018, -0.0039,  0.0033,  ...,  0.0059, -0.0068, -0.0069],\n",
      "        ...,\n",
      "        [-0.0086, -0.0054, -0.0006,  ..., -0.0041,  0.0102,  0.0047],\n",
      "        [-0.0043,  0.0086,  0.0020,  ...,  0.0005, -0.0075, -0.0024],\n",
      "        [-0.0082, -0.0082,  0.0078,  ..., -0.0058, -0.0072, -0.0050]],\n",
      "       device='cuda:0', requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([-5.3230e-03,  3.1211e-03, -8.4422e-04, -5.0879e-03,  9.4559e-03,\n",
      "        -9.5978e-03,  9.7345e-03, -1.4408e-03, -1.1375e-03, -3.9829e-03,\n",
      "         2.0466e-03,  4.6014e-03,  6.3707e-03,  8.7833e-03, -4.6859e-03,\n",
      "        -8.1269e-03, -1.3534e-03, -9.0286e-03, -9.5572e-03, -9.9541e-04,\n",
      "         3.2569e-03,  4.5758e-03, -8.0992e-03,  9.5596e-03,  1.9145e-03,\n",
      "         5.4675e-03, -1.3981e-03,  4.1981e-03,  8.8588e-03, -5.5976e-03,\n",
      "        -9.5791e-03, -8.0216e-03, -3.2328e-03, -7.0397e-03, -4.7427e-03,\n",
      "         2.5941e-03,  9.9955e-03,  5.6643e-03,  6.2606e-03, -6.7572e-03,\n",
      "         9.6156e-03, -1.4546e-03,  8.7670e-03,  5.4371e-03,  1.0326e-03,\n",
      "         9.8003e-03,  8.9032e-03, -1.4900e-05, -4.7388e-03, -8.7160e-03,\n",
      "        -7.5774e-05,  1.8449e-03, -2.6617e-03, -7.6649e-04, -9.4985e-03,\n",
      "         2.7917e-03, -6.1893e-03,  5.1694e-04, -1.9408e-03,  6.9423e-03,\n",
      "        -4.2699e-03,  3.5148e-03, -7.8502e-03,  5.8915e-03, -3.0433e-03,\n",
      "         8.3464e-05,  9.8500e-03, -7.0677e-04,  1.7415e-04, -2.1525e-03,\n",
      "         7.0959e-03, -8.2384e-04,  3.4531e-03,  5.0115e-03,  4.1892e-03,\n",
      "         6.8496e-03,  6.3136e-03,  7.3391e-03,  4.0054e-03,  6.8728e-05,\n",
      "         8.8970e-03, -8.5744e-03,  1.6070e-03, -8.6623e-04,  1.0121e-02,\n",
      "         2.8686e-03, -9.6245e-03, -8.0359e-03,  6.3255e-03, -1.4020e-03,\n",
      "        -1.5557e-03,  9.6665e-04,  4.2115e-03,  5.2187e-04,  6.8481e-04,\n",
      "         9.5811e-03,  4.6067e-03,  5.6818e-03,  6.0081e-03, -1.5273e-03,\n",
      "        -2.8688e-04, -9.3971e-03, -8.1499e-03,  2.3396e-03,  1.5670e-04,\n",
      "         1.4873e-03,  1.6849e-03,  9.5879e-03, -8.8051e-03, -6.0515e-03,\n",
      "         1.9446e-03, -6.7083e-03, -8.0252e-04,  3.3243e-03, -2.0455e-03,\n",
      "         7.7260e-03, -6.1639e-03, -1.5923e-03,  8.1961e-03, -3.1509e-03,\n",
      "         5.3824e-03, -6.5283e-03,  9.9530e-03, -7.7618e-03, -6.1090e-03,\n",
      "         7.4740e-03, -1.3899e-03, -9.7222e-03], device='cuda:0',\n",
      "       requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-0.0244, -0.0106, -0.0055,  ...,  0.0398, -0.0063, -0.0533],\n",
      "        [ 0.0752,  0.0853,  0.0820,  ...,  0.0204, -0.0878, -0.0264],\n",
      "        [-0.0679, -0.0404,  0.0238,  ..., -0.0824, -0.0145, -0.0003],\n",
      "        ...,\n",
      "        [-0.0438, -0.0266, -0.0722,  ..., -0.0801,  0.0561,  0.0149],\n",
      "        [-0.0583,  0.0035, -0.0866,  ..., -0.0708,  0.0593,  0.0561],\n",
      "        [-0.0018,  0.0219,  0.0358,  ...,  0.0223, -0.0170, -0.0371]],\n",
      "       device='cuda:0', requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([ 0.0150, -0.0186,  0.0150,  0.0385, -0.0663, -0.0434,  0.0075,  0.0427,\n",
      "        -0.0706,  0.0637], device='cuda:0', requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "for param in model.parameters():\n",
    "    print(param)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'conv1.weight': tensor([[[[[-1.6652e-02, -2.2975e-02,  1.9531e-04],\n",
       "            [ 2.5385e-02, -2.2959e-02, -1.2314e-02],\n",
       "            [-1.4160e-02, -1.8586e-03,  6.4033e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 2.1709e-03,  2.7222e-03,  3.5510e-03],\n",
       "            [ 6.0316e-03, -2.2517e-02,  3.3607e-03],\n",
       "            [ 1.6695e-02,  4.2553e-02, -2.8156e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 7.0681e-03, -1.2002e-02,  1.3895e-02],\n",
       "            [-2.9352e-02, -9.3097e-03, -5.0439e-03],\n",
       "            [ 4.7273e-03, -2.4532e-02,  1.0372e-02]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[ 4.8668e-03,  2.0931e-03,  7.6923e-03],\n",
       "            [-9.2889e-03, -1.1912e-02,  2.5961e-03],\n",
       "            [-7.6692e-03, -1.2963e-02,  4.3653e-03]]],\n",
       " \n",
       " \n",
       "          [[[-9.2796e-04,  8.0672e-03,  4.2809e-04],\n",
       "            [ 6.0073e-03,  9.2049e-03,  3.8828e-02],\n",
       "            [ 1.6068e-02,  1.1058e-03, -2.5562e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 9.2217e-03, -5.0106e-03, -6.9974e-03],\n",
       "            [ 9.8014e-03, -1.6241e-03,  1.1001e-02],\n",
       "            [-2.0466e-02, -3.3676e-03, -4.5282e-03]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-2.6881e-02, -9.7082e-03, -3.9761e-03],\n",
       "            [ 6.6590e-04, -7.3963e-03,  9.5708e-03],\n",
       "            [-9.4938e-03,  1.8030e-02,  1.9214e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 2.1393e-03,  1.1056e-02, -1.3760e-02],\n",
       "            [ 3.4600e-02,  1.7730e-03,  9.4388e-03],\n",
       "            [ 1.5448e-02, -1.5632e-02, -1.6828e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 2.1929e-03, -5.1399e-03, -2.3790e-02],\n",
       "            [-8.4550e-03,  1.9139e-02, -4.0301e-03],\n",
       "            [ 9.5540e-03,  2.8174e-02,  6.7652e-03]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-4.0700e-03,  4.6065e-04,  2.2577e-02],\n",
       "            [-7.5234e-03, -1.4119e-02, -1.3303e-02],\n",
       "            [ 3.1116e-04,  1.4258e-02,  1.2977e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.4834e-02,  4.9426e-03,  1.1928e-02],\n",
       "            [-6.7052e-03, -1.1831e-02,  5.2514e-03],\n",
       "            [-4.2482e-02, -7.8647e-03, -1.4218e-02]]],\n",
       " \n",
       " \n",
       "          [[[-8.3424e-03, -7.0247e-03, -4.7030e-03],\n",
       "            [ 1.4896e-02, -5.3225e-03,  1.1323e-02],\n",
       "            [ 2.7965e-03, -1.3497e-03, -1.3003e-02]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-4.2231e-03,  4.9511e-03,  1.7551e-02],\n",
       "            [ 2.7801e-02,  6.5378e-03,  2.3867e-02],\n",
       "            [-1.3652e-02,  2.2864e-03, -1.9444e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 1.9373e-02, -4.6635e-03, -1.9360e-02],\n",
       "            [ 1.6792e-02,  7.7016e-03,  2.6125e-02],\n",
       "            [-4.0033e-02, -1.0761e-02,  2.8261e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 1.9135e-02,  1.1607e-02,  6.4125e-03],\n",
       "            [ 3.1457e-02, -1.2191e-02, -3.4754e-02],\n",
       "            [ 3.3219e-03,  1.8991e-03, -2.0004e-02]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-2.3498e-02,  8.0592e-03, -1.8368e-02],\n",
       "            [-5.0842e-03, -8.9172e-03,  6.3292e-03],\n",
       "            [-1.2170e-02,  2.2580e-02,  1.9323e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.0111e-02, -1.8529e-02,  9.2907e-03],\n",
       "            [-3.2274e-03, -2.4153e-03, -1.4544e-04],\n",
       "            [ 1.1375e-02,  1.0150e-02, -1.3553e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 7.0373e-03,  1.9286e-02, -2.0108e-02],\n",
       "            [ 3.3312e-02,  6.0229e-03, -1.0761e-02],\n",
       "            [ 3.5367e-02, -2.8436e-02, -2.3368e-02]]]],\n",
       " \n",
       " \n",
       " \n",
       "         ...,\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-8.1843e-03, -2.5908e-02,  2.0439e-02],\n",
       "            [-7.6121e-03, -2.1403e-02, -8.5890e-03],\n",
       "            [ 8.6157e-03,  5.5030e-03, -4.3926e-03]]],\n",
       " \n",
       " \n",
       "          [[[-2.4904e-02,  2.4799e-02,  3.0888e-03],\n",
       "            [ 1.8908e-02, -7.7103e-03,  1.8519e-02],\n",
       "            [ 2.8565e-02,  1.5077e-03, -3.1009e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.7129e-02, -5.5274e-03,  2.0510e-04],\n",
       "            [ 1.3983e-02, -1.1021e-02,  1.1323e-02],\n",
       "            [-1.3245e-02, -1.7739e-02,  1.4229e-02]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[ 1.1729e-02,  1.9461e-02,  3.2485e-02],\n",
       "            [-3.6170e-03,  1.9356e-02,  1.5700e-02],\n",
       "            [-1.5682e-02, -1.7206e-03,  1.9711e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.1009e-02,  1.0090e-02,  3.3992e-03],\n",
       "            [-1.6689e-02, -4.4330e-02, -2.7226e-03],\n",
       "            [ 9.9930e-05, -1.2215e-02, -1.1349e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 9.0160e-03,  2.4521e-02,  5.2484e-04],\n",
       "            [-7.6661e-03,  1.6224e-03, -7.1766e-03],\n",
       "            [ 8.8503e-03, -1.0304e-02,  2.4501e-04]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-1.9703e-02, -2.3008e-02, -1.2450e-02],\n",
       "            [-5.3685e-03, -8.0234e-03,  5.6956e-03],\n",
       "            [-3.3180e-02,  4.5688e-03, -1.8126e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.5145e-02,  1.0309e-02,  7.7255e-04],\n",
       "            [ 2.8422e-03, -3.8337e-03,  5.8257e-03],\n",
       "            [-2.6666e-02,  3.7327e-02, -2.0397e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.3473e-02, -1.1163e-02,  1.3897e-02],\n",
       "            [-3.9547e-03, -6.4102e-03,  7.2545e-03],\n",
       "            [-3.2708e-04,  1.3834e-02, -1.3423e-02]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-1.7094e-02,  3.3743e-03,  2.0040e-02],\n",
       "            [ 2.4377e-03,  5.4402e-03, -2.2066e-02],\n",
       "            [-1.0437e-02,  1.1204e-02,  2.0456e-04]]],\n",
       " \n",
       " \n",
       "          [[[ 1.3700e-02, -1.6171e-03, -1.9126e-02],\n",
       "            [-3.8071e-02,  1.0903e-02, -7.4041e-03],\n",
       "            [ 1.3000e-02, -1.2129e-02,  1.3393e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 9.0568e-03,  6.8606e-03, -2.4280e-02],\n",
       "            [ 5.0935e-03, -8.2808e-03,  2.3902e-04],\n",
       "            [ 2.6941e-02,  2.6772e-03,  9.0850e-03]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[ 1.8270e-03, -8.3236e-03, -2.5054e-04],\n",
       "            [ 1.2092e-02,  1.2707e-02,  3.0889e-02],\n",
       "            [-2.1067e-02,  9.0260e-05,  9.3656e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 2.3230e-03,  2.5367e-02,  4.8213e-03],\n",
       "            [-5.4665e-03,  6.0833e-03, -4.9704e-03],\n",
       "            [ 3.2607e-02,  4.1543e-02, -3.0244e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 7.7684e-03,  1.5508e-02,  1.0407e-03],\n",
       "            [ 8.7677e-03, -2.3000e-02,  2.0225e-02],\n",
       "            [-8.3651e-03, -1.0295e-02, -6.6316e-03]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-5.7536e-03, -2.3419e-02,  2.8286e-03],\n",
       "            [ 3.5407e-03, -2.4306e-02, -1.1785e-02],\n",
       "            [-2.5490e-03, -1.5586e-03, -8.3251e-04]]],\n",
       " \n",
       " \n",
       "          [[[-1.3981e-02, -8.3511e-04,  2.5361e-02],\n",
       "            [-4.0168e-03,  7.8830e-04,  2.8949e-02],\n",
       "            [ 2.1428e-02,  3.5264e-03,  6.3982e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 7.6870e-03,  4.0428e-03, -2.4760e-02],\n",
       "            [ 1.0954e-02, -4.3701e-03, -2.9109e-02],\n",
       "            [-8.1288e-03, -2.6292e-03, -6.9497e-03]]]]], device='cuda:0'),\n",
       " 'conv1.bias': tensor([[-0.0162, -0.0111, -0.0006,  ...,  0.0011, -0.0025,  0.0082],\n",
       "         [ 0.0005, -0.0105, -0.0007,  ...,  0.0117, -0.0184,  0.0138],\n",
       "         [-0.0250,  0.0234,  0.0103,  ...,  0.0009,  0.0009,  0.0072],\n",
       "         ...,\n",
       "         [ 0.0034, -0.0133,  0.0050,  ...,  0.0255, -0.0023, -0.0078],\n",
       "         [-0.0323,  0.0143,  0.0226,  ..., -0.0097,  0.0123,  0.0014],\n",
       "         [ 0.0036, -0.0094, -0.0112,  ...,  0.0107, -0.0121, -0.0028]],\n",
       "        device='cuda:0'),\n",
       " 'conv2.weight': tensor([[[[[ 3.0730e-03, -2.4890e-02, -1.2325e-02],\n",
       "            [-2.0468e-02, -1.6916e-02, -1.1415e-02],\n",
       "            [-9.8524e-03, -2.8356e-02, -2.9868e-02]],\n",
       " \n",
       "           [[-2.1364e-02, -1.0517e-02, -2.2693e-02],\n",
       "            [-2.4282e-02, -1.6433e-02, -3.0771e-02],\n",
       "            [-3.0342e-02, -1.8309e-02, -2.9933e-02]],\n",
       " \n",
       "           [[-2.4358e-02, -1.5849e-02, -1.8873e-02],\n",
       "            [-6.9543e-03, -4.4690e-03, -2.9589e-02],\n",
       "            [-2.8684e-02, -1.4762e-02, -1.4690e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-6.8019e-04, -6.3734e-03, -8.7640e-03],\n",
       "            [ 1.6306e-03, -1.5812e-02, -4.3925e-03],\n",
       "            [-9.5445e-03, -1.0529e-02, -9.8536e-03]],\n",
       " \n",
       "           [[-3.2482e-02, -1.2923e-02, -1.4186e-02],\n",
       "            [-4.1521e-03, -1.8549e-02, -1.3925e-02],\n",
       "            [-2.5084e-02, -1.2564e-02, -2.2636e-02]],\n",
       " \n",
       "           [[-8.9581e-03, -9.7553e-03, -2.3417e-03],\n",
       "            [-1.3273e-02, -4.7430e-03, -2.2406e-02],\n",
       "            [-1.3902e-02, -7.9354e-03, -2.9658e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 1.0363e-02, -6.7375e-03,  1.0538e-03],\n",
       "            [-1.5655e-02, -2.1905e-02,  3.3508e-03],\n",
       "            [-4.8086e-03,  1.2671e-03, -2.2850e-03]],\n",
       " \n",
       "           [[ 5.7121e-03,  6.9669e-03, -5.8917e-03],\n",
       "            [ 2.2283e-03, -8.7533e-03, -1.1485e-02],\n",
       "            [ 4.2973e-03,  3.2817e-03,  6.2982e-03]],\n",
       " \n",
       "           [[ 8.6416e-03, -7.5832e-03, -9.8566e-03],\n",
       "            [ 7.2141e-03, -2.9093e-03, -1.5504e-02],\n",
       "            [-9.3326e-03, -3.4596e-03,  1.5121e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.8353e-03, -1.3265e-02, -7.7008e-03],\n",
       "            [ 2.7088e-03,  3.1760e-04,  9.1666e-03],\n",
       "            [-1.3500e-02, -3.2598e-03,  4.4771e-03]],\n",
       " \n",
       "           [[-7.7261e-03, -4.4223e-03,  5.9535e-03],\n",
       "            [-3.1848e-03,  1.0236e-02, -5.4610e-03],\n",
       "            [ 1.2943e-02,  1.6069e-02, -8.6001e-03]],\n",
       " \n",
       "           [[ 1.3856e-03, -1.1102e-03,  3.5356e-03],\n",
       "            [ 4.6940e-04, -3.5052e-03, -8.1128e-03],\n",
       "            [ 6.1435e-04,  7.7807e-03,  2.7160e-03]]],\n",
       " \n",
       " \n",
       "          [[[-8.3658e-03, -1.5346e-02, -1.1856e-02],\n",
       "            [-1.7005e-02, -2.0792e-02, -2.5550e-02],\n",
       "            [-2.5601e-02, -9.6715e-03, -1.8328e-02]],\n",
       " \n",
       "           [[-2.3792e-03, -7.1213e-03, -1.9534e-02],\n",
       "            [-1.7613e-02, -2.4310e-02, -1.8925e-02],\n",
       "            [-1.4016e-02, -1.1434e-02, -6.2763e-03]],\n",
       " \n",
       "           [[-1.2591e-02,  1.1056e-03, -6.8569e-03],\n",
       "            [-1.4024e-02, -2.2343e-02, -9.6856e-03],\n",
       "            [ 4.0197e-03, -1.4607e-02, -2.2973e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-5.5233e-03, -1.4175e-02, -6.9893e-03],\n",
       "            [-1.1738e-02, -9.9921e-03, -2.3020e-03],\n",
       "            [-6.8148e-03, -1.9962e-03,  1.7750e-03]],\n",
       " \n",
       "           [[-4.8195e-03, -6.3706e-03, -1.7380e-02],\n",
       "            [-1.3988e-02, -1.5155e-02, -2.1725e-02],\n",
       "            [-5.6177e-03, -7.8411e-03, -8.3893e-03]],\n",
       " \n",
       "           [[-6.1376e-03, -8.4482e-04, -2.7049e-03],\n",
       "            [-8.7489e-03, -1.9494e-02, -1.2324e-02],\n",
       "            [-1.2818e-03, -6.2487e-03, -6.4332e-03]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[ 1.3029e-03,  8.1606e-03, -7.0689e-03],\n",
       "            [ 7.0847e-03,  1.5219e-02,  4.2621e-03],\n",
       "            [-2.3104e-03,  5.8669e-03, -2.2329e-03]],\n",
       " \n",
       "           [[ 1.3356e-02,  9.4446e-03,  4.4198e-03],\n",
       "            [ 7.8391e-03,  2.2959e-03,  1.3148e-02],\n",
       "            [ 1.2221e-02, -5.9384e-03,  7.6757e-03]],\n",
       " \n",
       "           [[ 6.0576e-03,  1.1265e-02,  7.7541e-03],\n",
       "            [-5.4630e-03,  1.4956e-03,  9.2534e-03],\n",
       "            [ 1.3805e-02, -2.9113e-03, -1.6178e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.0727e-03,  5.4815e-03,  6.1781e-03],\n",
       "            [-3.5164e-03,  6.5823e-04, -1.2303e-02],\n",
       "            [ 4.2402e-03,  4.4931e-03, -7.0074e-03]],\n",
       " \n",
       "           [[ 1.1327e-02,  1.5607e-02,  4.0116e-03],\n",
       "            [ 1.2961e-03, -3.8904e-03,  8.0646e-03],\n",
       "            [-1.0875e-03, -2.0722e-02,  1.2624e-03]],\n",
       " \n",
       "           [[ 6.2738e-03,  5.8141e-03,  5.8049e-03],\n",
       "            [ 3.9719e-03, -2.1144e-03,  6.1054e-03],\n",
       "            [ 2.7190e-03, -4.1102e-03,  7.7675e-04]]],\n",
       " \n",
       " \n",
       "          [[[-5.0158e-03, -1.2162e-02, -6.6592e-03],\n",
       "            [ 4.7244e-04, -1.4373e-03, -2.0920e-02],\n",
       "            [ 7.1782e-03, -1.1650e-02,  1.1314e-02]],\n",
       " \n",
       "           [[-7.5319e-03, -2.6262e-03, -2.2832e-03],\n",
       "            [-5.4604e-03, -7.8393e-03, -1.6143e-02],\n",
       "            [ 4.3753e-04,  2.8509e-03, -6.6886e-03]],\n",
       " \n",
       "           [[-7.1595e-03, -1.0709e-02, -1.3357e-03],\n",
       "            [ 2.2635e-03, -1.5594e-02, -7.2155e-03],\n",
       "            [ 1.5190e-02,  6.0235e-04, -1.9403e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 6.6572e-03, -2.2231e-03, -1.6868e-03],\n",
       "            [ 6.3627e-03,  6.9596e-04,  6.0099e-03],\n",
       "            [-5.0518e-03, -1.2405e-02,  2.2555e-03]],\n",
       " \n",
       "           [[-8.2779e-03,  2.2690e-03, -2.0877e-02],\n",
       "            [ 1.6578e-03, -8.3872e-03,  7.4535e-04],\n",
       "            [-7.2402e-04,  1.0000e-03,  1.4927e-02]],\n",
       " \n",
       "           [[-5.6660e-03, -4.1564e-03, -1.8084e-03],\n",
       "            [ 1.3408e-03, -4.7343e-03, -1.5874e-02],\n",
       "            [ 7.0109e-03,  2.9490e-03,  4.7196e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 2.4785e-03,  9.6384e-03, -2.8784e-02],\n",
       "            [-1.7531e-02,  8.3649e-04,  1.2180e-02],\n",
       "            [ 2.9646e-03,  1.7013e-02, -3.2355e-03]],\n",
       " \n",
       "           [[ 1.3212e-03,  7.7659e-03, -2.1624e-03],\n",
       "            [-1.7751e-05, -7.1681e-03, -8.2588e-03],\n",
       "            [ 1.1545e-02,  1.0329e-02,  2.0599e-03]],\n",
       " \n",
       "           [[-6.3528e-03, -2.8322e-02,  1.7491e-02],\n",
       "            [ 4.6166e-03, -5.0771e-03, -2.8369e-02],\n",
       "            [ 4.7991e-03,  1.1128e-02,  9.1124e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.6866e-03, -1.1756e-02,  1.1136e-02],\n",
       "            [ 4.4958e-03, -3.4539e-03,  6.8649e-03],\n",
       "            [-9.3391e-03, -1.3224e-03, -1.2070e-02]],\n",
       " \n",
       "           [[ 3.8268e-03,  2.2354e-02,  9.1852e-04],\n",
       "            [-6.2498e-03,  3.6643e-03,  3.7118e-03],\n",
       "            [-5.2775e-03,  8.6279e-03,  3.8593e-03]],\n",
       " \n",
       "           [[ 3.5888e-03,  7.9377e-03,  7.2320e-03],\n",
       "            [-1.4060e-03, -4.2298e-03, -8.9105e-03],\n",
       "            [ 4.1519e-03,  7.8844e-03,  4.8114e-03]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-6.0205e-04,  1.9908e-02,  1.1630e-02],\n",
       "            [ 1.5035e-02,  1.2113e-02,  2.3543e-02],\n",
       "            [ 1.2792e-02,  5.6260e-03,  1.7471e-02]],\n",
       " \n",
       "           [[ 1.4662e-02,  1.4119e-02,  8.8687e-03],\n",
       "            [ 2.1147e-02,  1.0077e-02,  1.0267e-02],\n",
       "            [ 1.1556e-02,  1.6680e-02, -4.2962e-03]],\n",
       " \n",
       "           [[ 2.4570e-02,  2.3951e-02,  2.3056e-05],\n",
       "            [ 2.7605e-03,  2.0880e-02,  1.4921e-02],\n",
       "            [ 2.1450e-02,  2.0968e-02, -8.6795e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 1.5481e-02,  4.5844e-03,  9.4296e-03],\n",
       "            [ 5.7479e-03,  2.1795e-03, -6.4201e-04],\n",
       "            [ 1.2684e-02, -8.9133e-03,  1.8599e-02]],\n",
       " \n",
       "           [[-8.4927e-04,  4.0223e-03,  2.1723e-02],\n",
       "            [ 8.9688e-03,  1.9875e-02,  1.2017e-02],\n",
       "            [ 2.7487e-02, -9.6626e-03,  1.8683e-02]],\n",
       " \n",
       "           [[ 8.0839e-03,  9.3684e-03,  2.2500e-03],\n",
       "            [ 7.6510e-03,  7.7052e-03,  7.4984e-03],\n",
       "            [ 6.5674e-03,  8.0957e-03,  1.6001e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 1.4464e-03,  2.2115e-03,  1.0468e-02],\n",
       "            [ 8.1808e-03, -1.3305e-02,  9.7301e-03],\n",
       "            [ 3.0892e-03,  1.5045e-02, -3.6762e-03]],\n",
       " \n",
       "           [[ 6.4295e-03,  1.2924e-02, -3.8864e-03],\n",
       "            [-5.0039e-03,  6.2470e-03, -6.8365e-03],\n",
       "            [-2.1867e-04,  7.7893e-03,  1.2418e-02]],\n",
       " \n",
       "           [[ 4.9479e-03, -5.2812e-03,  9.8762e-04],\n",
       "            [ 9.0913e-03,  2.2592e-03,  2.6459e-03],\n",
       "            [ 1.2722e-04, -1.1726e-02,  1.8688e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 7.9418e-03, -1.5031e-02, -5.9884e-04],\n",
       "            [-1.4858e-03, -9.5558e-04,  1.2342e-02],\n",
       "            [ 2.6953e-03, -4.8014e-03, -3.0186e-03]],\n",
       " \n",
       "           [[ 1.5005e-02,  7.5051e-03,  5.2993e-03],\n",
       "            [ 1.2898e-03,  8.6432e-03, -2.8607e-02],\n",
       "            [ 2.5212e-04,  3.0458e-02, -2.6272e-03]],\n",
       " \n",
       "           [[ 9.1546e-03,  8.6584e-03, -2.0269e-03],\n",
       "            [ 4.4659e-03,  6.1750e-03, -2.8844e-03],\n",
       "            [ 1.0696e-03,  5.8112e-03,  9.3345e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 6.8601e-04,  6.3889e-03, -1.4806e-03],\n",
       "            [ 1.9948e-02, -9.1333e-04,  8.3264e-04],\n",
       "            [ 6.4929e-03, -5.3051e-03,  1.4967e-02]],\n",
       " \n",
       "           [[-2.3941e-03,  1.9838e-03,  1.3252e-02],\n",
       "            [ 6.3840e-04,  1.3196e-02,  4.0120e-03],\n",
       "            [ 3.1659e-03,  1.8219e-02,  1.2273e-02]],\n",
       " \n",
       "           [[-3.9154e-03,  9.1355e-03,  1.5629e-03],\n",
       "            [ 4.0868e-03,  1.3046e-02,  6.8529e-03],\n",
       "            [-9.3882e-03,  1.3246e-02,  1.0755e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.9590e-03,  1.0721e-02, -3.5016e-04],\n",
       "            [ 2.8563e-03,  6.2428e-03,  2.6555e-03],\n",
       "            [ 8.0119e-03,  5.3724e-03, -2.0522e-03]],\n",
       " \n",
       "           [[ 1.5036e-02, -1.6108e-03,  4.0009e-03],\n",
       "            [ 2.8910e-03, -8.6156e-04,  1.8680e-02],\n",
       "            [ 8.8895e-03,  1.6461e-02,  1.3200e-03]],\n",
       " \n",
       "           [[-3.7821e-03,  9.3575e-04, -1.8557e-03],\n",
       "            [ 2.0519e-03, -2.7137e-03,  4.3127e-03],\n",
       "            [-3.0105e-03,  1.0031e-02,  8.8200e-03]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-7.0375e-03, -5.9731e-03,  1.3146e-02],\n",
       "            [-7.8030e-04, -2.4201e-03, -1.9666e-03],\n",
       "            [-1.2818e-02, -2.5406e-04, -8.5585e-03]],\n",
       " \n",
       "           [[-4.7869e-03, -8.5239e-03, -1.8915e-03],\n",
       "            [-9.6075e-03,  2.1396e-03, -1.6341e-03],\n",
       "            [-6.3538e-03,  6.7600e-04, -6.4958e-03]],\n",
       " \n",
       "           [[ 4.3501e-03,  6.8479e-03, -1.1757e-02],\n",
       "            [-9.8394e-03,  1.1029e-03,  7.4635e-03],\n",
       "            [ 4.4795e-03,  8.5821e-03, -3.4606e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-6.1850e-03,  3.5213e-03,  6.5784e-03],\n",
       "            [-2.4762e-03, -3.8353e-03, -3.7063e-03],\n",
       "            [ 7.5290e-03,  7.0813e-04, -8.2252e-03]],\n",
       " \n",
       "           [[ 3.6745e-03, -1.3438e-02, -7.1533e-03],\n",
       "            [-5.1337e-03, -7.0800e-03, -5.1736e-03],\n",
       "            [-5.0745e-04, -1.1321e-02, -2.0504e-03]],\n",
       " \n",
       "           [[ 1.2978e-03,  1.3022e-03, -5.1487e-03],\n",
       "            [-6.2732e-03, -3.1590e-03,  4.6786e-03],\n",
       "            [-7.6378e-03, -2.1105e-03, -5.4404e-03]]],\n",
       " \n",
       " \n",
       "          [[[-9.4658e-03, -8.6112e-03, -9.3819e-03],\n",
       "            [-5.7105e-03, -7.9600e-03, -1.7263e-03],\n",
       "            [ 5.5935e-03, -7.4578e-03,  2.8307e-03]],\n",
       " \n",
       "           [[ 3.6013e-05,  3.6909e-03, -1.2593e-02],\n",
       "            [-2.6929e-03, -2.4020e-03, -1.4745e-02],\n",
       "            [-2.5663e-03, -9.7284e-04, -4.5484e-03]],\n",
       " \n",
       "           [[ 6.7365e-04, -1.5376e-02, -1.4209e-02],\n",
       "            [-2.1198e-03, -4.5921e-03, -1.8000e-02],\n",
       "            [-1.4946e-02, -3.5640e-03,  1.0054e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 1.3419e-02, -9.2282e-03, -3.4893e-03],\n",
       "            [ 2.3933e-04, -7.1516e-03,  4.6803e-03],\n",
       "            [-3.0322e-03, -1.1138e-02, -5.0506e-04]],\n",
       " \n",
       "           [[-2.0303e-02, -6.0179e-03,  1.8111e-03],\n",
       "            [-7.1650e-03, -2.9188e-03,  9.2878e-03],\n",
       "            [ 6.9947e-04,  1.0597e-02,  5.3455e-03]],\n",
       " \n",
       "           [[ 2.6131e-04, -3.3793e-05, -2.9704e-03],\n",
       "            [-4.2516e-03, -2.8968e-05, -8.5048e-03],\n",
       "            [-4.5677e-03,  3.3091e-03,  2.4250e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 1.3528e-02, -1.7258e-02, -1.3862e-03],\n",
       "            [-4.8493e-03, -6.5021e-03,  9.4963e-03],\n",
       "            [-8.4563e-03,  2.0767e-02,  1.0482e-02]],\n",
       " \n",
       "           [[ 5.1152e-03,  1.7524e-03, -7.0389e-03],\n",
       "            [-6.7651e-03, -3.7812e-03, -8.6203e-03],\n",
       "            [ 1.5002e-03,  7.4817e-03,  1.5098e-02]],\n",
       " \n",
       "           [[-2.0526e-02,  1.8126e-03,  1.9487e-03],\n",
       "            [ 1.7752e-02, -2.1877e-03, -1.6935e-02],\n",
       "            [-5.0958e-04,  3.0314e-03,  2.4242e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-8.4036e-04, -3.6430e-03, -3.3423e-03],\n",
       "            [-3.0306e-03, -7.0898e-04,  1.0511e-02],\n",
       "            [ 1.7542e-03, -5.9962e-04,  1.0659e-02]],\n",
       " \n",
       "           [[ 2.5921e-03, -3.1498e-03,  2.6501e-02],\n",
       "            [-3.8152e-03,  1.1981e-03, -3.3704e-03],\n",
       "            [-4.4323e-04,  6.0158e-03, -9.6770e-03]],\n",
       " \n",
       "           [[-1.1515e-03, -2.8627e-03,  1.9279e-03],\n",
       "            [ 2.8486e-03, -2.2285e-03, -9.5412e-03],\n",
       "            [-4.5255e-03,  7.0836e-03,  1.1640e-02]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[ 2.1867e-03,  1.5008e-02,  1.5212e-02],\n",
       "            [ 2.4449e-02,  9.0459e-03,  1.4918e-02],\n",
       "            [-1.2502e-03, -1.6144e-02,  3.1699e-03]],\n",
       " \n",
       "           [[ 1.2841e-02,  1.3269e-02,  1.3706e-02],\n",
       "            [ 1.6946e-03,  1.5737e-02,  7.0122e-03],\n",
       "            [-4.5302e-03, -7.8719e-03, -1.3459e-02]],\n",
       " \n",
       "           [[ 1.1551e-02,  1.5054e-02,  2.3829e-03],\n",
       "            [-1.3162e-03,  1.2648e-02,  6.5400e-03],\n",
       "            [ 8.0030e-03, -1.3412e-02, -1.3476e-04]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 9.2451e-03,  5.2481e-03,  8.7361e-03],\n",
       "            [ 3.7308e-03, -4.6749e-05,  2.4103e-03],\n",
       "            [ 1.9548e-03,  3.5664e-03, -4.1559e-03]],\n",
       " \n",
       "           [[ 2.7194e-03,  7.7213e-03,  2.5833e-02],\n",
       "            [ 1.4263e-02, -5.0782e-03,  2.0793e-02],\n",
       "            [-7.8842e-03, -3.7313e-03, -4.4826e-03]],\n",
       " \n",
       "           [[ 3.5611e-03,  1.0139e-02,  4.2338e-03],\n",
       "            [-2.3665e-03,  3.4680e-03,  1.1440e-02],\n",
       "            [-4.5074e-03, -7.9678e-03, -2.4344e-03]]],\n",
       " \n",
       " \n",
       "          [[[-5.0122e-03,  8.1858e-04,  1.4083e-02],\n",
       "            [ 1.5162e-05, -6.4715e-03,  1.4379e-03],\n",
       "            [-8.8809e-03,  1.0624e-02, -8.1843e-03]],\n",
       " \n",
       "           [[-6.0179e-03,  5.1558e-03, -1.3236e-03],\n",
       "            [-5.7526e-03,  1.5404e-03, -5.4450e-03],\n",
       "            [-3.7715e-04, -2.0126e-03,  5.4982e-03]],\n",
       " \n",
       "           [[ 5.4488e-03,  2.1275e-02, -8.8682e-03],\n",
       "            [ 4.0211e-03, -4.0317e-03, -3.3759e-04],\n",
       "            [ 4.8645e-03, -2.5775e-03,  2.0777e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 1.3891e-02,  7.1964e-03, -6.0635e-04],\n",
       "            [ 1.7643e-03,  3.0493e-03,  1.0427e-02],\n",
       "            [-6.5914e-04, -2.3012e-03,  1.5230e-03]],\n",
       " \n",
       "           [[-6.2575e-04, -1.3734e-02,  1.5897e-03],\n",
       "            [-1.3519e-02,  1.3268e-02,  4.1617e-03],\n",
       "            [-3.6795e-03,  1.4166e-02, -1.2312e-02]],\n",
       " \n",
       "           [[-8.4208e-03,  1.8405e-03, -3.8895e-03],\n",
       "            [ 3.3187e-03, -6.0134e-03, -2.4714e-03],\n",
       "            [-5.6897e-03,  1.8374e-04,  3.1190e-04]]],\n",
       " \n",
       " \n",
       "          [[[-5.8944e-03,  1.4110e-02, -4.5760e-03],\n",
       "            [ 1.4281e-02,  7.5068e-03,  5.3258e-04],\n",
       "            [ 1.6691e-02,  4.7528e-03,  1.7107e-02]],\n",
       " \n",
       "           [[ 2.6984e-04, -3.7285e-03,  1.0389e-02],\n",
       "            [ 1.1082e-02,  7.6261e-03,  3.1418e-03],\n",
       "            [ 2.8875e-02,  1.9037e-02,  8.0851e-03]],\n",
       " \n",
       "           [[ 5.8787e-04, -1.0733e-02, -2.0443e-03],\n",
       "            [ 8.8734e-03,  1.1207e-02, -5.4590e-03],\n",
       "            [ 8.6588e-03,  9.0354e-03,  2.1327e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.9510e-04, -5.5097e-03,  1.1268e-02],\n",
       "            [ 4.6934e-03,  1.0582e-02,  4.2402e-03],\n",
       "            [-1.0340e-03, -1.0672e-04, -3.4624e-03]],\n",
       " \n",
       "           [[-1.2542e-03, -7.5566e-03, -8.8580e-03],\n",
       "            [ 1.6081e-02, -7.9596e-03,  1.3792e-02],\n",
       "            [ 7.6254e-03,  3.9538e-02,  2.0022e-02]],\n",
       " \n",
       "           [[-7.6189e-03, -2.0555e-03,  6.2537e-04],\n",
       "            [ 6.7006e-03,  1.7283e-03,  3.0639e-04],\n",
       "            [ 1.7364e-02,  1.0374e-02,  1.7061e-02]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-1.1911e-02,  8.2663e-03, -8.4632e-03],\n",
       "            [-6.3491e-03,  1.6852e-03, -9.2396e-03],\n",
       "            [-8.6432e-03, -1.8442e-02,  3.6131e-03]],\n",
       " \n",
       "           [[ 3.7331e-03, -9.1073e-03,  1.7430e-03],\n",
       "            [-5.3076e-03, -1.6231e-02,  4.1240e-04],\n",
       "            [-7.2575e-04, -1.5138e-02, -1.0786e-02]],\n",
       " \n",
       "           [[ 1.0683e-03, -4.9705e-03,  6.9949e-03],\n",
       "            [-6.7406e-03, -6.6321e-03,  5.3634e-03],\n",
       "            [-1.3134e-02, -1.0272e-02, -1.0270e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.3817e-03, -6.4348e-03, -1.1167e-02],\n",
       "            [-9.2879e-03, -1.5280e-03, -1.8405e-03],\n",
       "            [-1.3225e-02, -5.9530e-03, -1.0757e-02]],\n",
       " \n",
       "           [[-1.1197e-02, -9.0014e-03, -5.0321e-04],\n",
       "            [-1.2739e-03, -2.2517e-02,  1.8146e-02],\n",
       "            [-8.2286e-03, -5.1718e-04, -3.5921e-03]],\n",
       " \n",
       "           [[ 3.7659e-03, -3.3008e-03, -7.2478e-04],\n",
       "            [-3.7360e-03, -1.1012e-02, -6.5364e-03],\n",
       "            [-5.3482e-03, -1.7003e-02, -9.9247e-03]]],\n",
       " \n",
       " \n",
       "          [[[-3.6458e-02,  5.6881e-03, -1.1582e-02],\n",
       "            [-7.1437e-03, -9.0720e-04,  4.4770e-03],\n",
       "            [ 9.5810e-03,  8.6882e-03,  1.0409e-02]],\n",
       " \n",
       "           [[ 1.0541e-02, -1.7995e-02, -1.5810e-03],\n",
       "            [ 1.1837e-02, -1.6842e-02, -1.2198e-04],\n",
       "            [ 1.0846e-02, -6.3608e-03, -7.5681e-05]],\n",
       " \n",
       "           [[ 3.1424e-03,  3.0978e-03,  6.0729e-03],\n",
       "            [-1.0197e-02, -7.9037e-03, -2.9961e-03],\n",
       "            [ 2.4360e-03,  1.0174e-02, -1.3039e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-2.0833e-02,  5.5670e-03, -7.6647e-03],\n",
       "            [-5.8758e-03,  6.1475e-03,  2.1956e-03],\n",
       "            [-2.1868e-02, -1.3689e-02, -9.6410e-03]],\n",
       " \n",
       "           [[-3.0386e-03, -1.5291e-02, -1.7933e-02],\n",
       "            [-9.5209e-03, -2.2636e-04, -4.4796e-03],\n",
       "            [ 2.6175e-02,  5.0363e-03, -1.3335e-02]],\n",
       " \n",
       "           [[ 3.2362e-04, -6.3070e-03, -5.9073e-03],\n",
       "            [ 2.0257e-03, -1.2151e-02, -6.6760e-03],\n",
       "            [ 1.2296e-02,  4.1050e-04, -5.3289e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 1.2417e-04,  1.0545e-02, -6.7183e-03],\n",
       "            [ 2.5361e-02,  4.0590e-03,  9.8807e-03],\n",
       "            [-5.5509e-03,  2.0680e-02, -6.2005e-03]],\n",
       " \n",
       "           [[ 6.0626e-03,  4.2464e-03,  1.9175e-03],\n",
       "            [ 6.0054e-03,  6.6450e-03,  1.4369e-02],\n",
       "            [-5.3424e-03, -1.0914e-03,  9.8609e-03]],\n",
       " \n",
       "           [[ 1.2836e-02,  4.3401e-03,  4.2322e-03],\n",
       "            [-9.7475e-03,  1.2843e-03,  7.7501e-03],\n",
       "            [ 2.4433e-02,  2.0401e-03,  4.1851e-04]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 9.8814e-03,  1.0879e-02,  7.9779e-03],\n",
       "            [-8.0392e-03,  1.6642e-03, -9.6601e-03],\n",
       "            [ 6.7883e-03,  8.9648e-03,  1.1463e-02]],\n",
       " \n",
       "           [[ 2.0115e-02,  3.4275e-03,  1.7117e-02],\n",
       "            [ 1.4522e-02,  3.6797e-02,  1.0411e-02],\n",
       "            [-1.9274e-04, -6.7919e-03,  1.1091e-02]],\n",
       " \n",
       "           [[-1.1682e-04,  5.1326e-03, -8.6268e-03],\n",
       "            [ 7.4306e-03,  4.9966e-03,  1.6057e-02],\n",
       "            [-6.1433e-03,  1.3681e-04,  2.6383e-05]]]],\n",
       " \n",
       " \n",
       " \n",
       "         ...,\n",
       " \n",
       " \n",
       " \n",
       "         [[[[ 8.3781e-03, -1.4320e-02,  1.0562e-03],\n",
       "            [-4.7683e-03, -1.9312e-02, -1.2858e-02],\n",
       "            [-1.9621e-02, -2.1851e-02, -2.2426e-03]],\n",
       " \n",
       "           [[-1.3888e-02,  5.4064e-03, -2.4987e-02],\n",
       "            [-2.0253e-02, -8.2049e-03, -2.9169e-02],\n",
       "            [-3.1986e-02, -1.1377e-02, -1.2378e-02]],\n",
       " \n",
       "           [[-8.3370e-03,  4.4071e-04, -4.5519e-03],\n",
       "            [-1.6535e-03, -2.4371e-02, -5.8397e-03],\n",
       "            [-2.6838e-02, -1.3097e-02, -8.8673e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 1.1426e-02, -4.7712e-03, -1.1970e-03],\n",
       "            [-5.5966e-03, -1.0843e-02,  8.5913e-03],\n",
       "            [ 5.5685e-03, -1.2251e-02,  2.7885e-03]],\n",
       " \n",
       "           [[-1.6164e-02, -9.3918e-03, -2.0831e-02],\n",
       "            [-3.4964e-02, -7.7173e-03, -1.4889e-02],\n",
       "            [-1.6988e-02, -9.3636e-03, -1.9245e-02]],\n",
       " \n",
       "           [[-6.8820e-03, -6.9086e-03, -1.2483e-02],\n",
       "            [-3.4990e-03, -8.3268e-03, -1.6587e-02],\n",
       "            [-3.0983e-02, -1.0508e-02, -1.9901e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 1.6792e-03,  1.3199e-02,  7.5025e-03],\n",
       "            [-6.9142e-03,  4.2362e-02,  7.3266e-03],\n",
       "            [ 4.2233e-03,  1.2737e-02,  6.1593e-03]],\n",
       " \n",
       "           [[ 1.6766e-02,  7.6832e-03,  1.9180e-02],\n",
       "            [ 2.4457e-02,  7.8239e-03,  2.9806e-02],\n",
       "            [ 2.7630e-02, -5.7175e-03,  2.3393e-02]],\n",
       " \n",
       "           [[ 1.3766e-02,  1.0197e-02,  2.2118e-02],\n",
       "            [ 2.0041e-02,  1.2011e-03,  2.2649e-02],\n",
       "            [ 2.0495e-03,  1.1207e-02,  6.7364e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 4.4552e-03,  1.1681e-03,  3.0871e-03],\n",
       "            [-4.6654e-03,  7.8024e-03, -5.9157e-04],\n",
       "            [-9.1663e-03,  2.3227e-02, -5.3320e-03]],\n",
       " \n",
       "           [[ 8.1281e-03,  4.2601e-02,  8.7673e-03],\n",
       "            [ 9.4691e-03,  1.3581e-02,  1.9127e-02],\n",
       "            [ 8.4107e-03,  1.3905e-02, -9.9739e-03]],\n",
       " \n",
       "           [[ 1.5001e-02,  6.1620e-03,  1.3385e-02],\n",
       "            [ 1.2983e-02,  8.7540e-03,  1.2166e-02],\n",
       "            [ 4.7086e-03, -1.6530e-03,  1.1324e-02]]],\n",
       " \n",
       " \n",
       "          [[[-7.3761e-03, -2.0316e-02, -1.7004e-02],\n",
       "            [-1.3624e-02, -1.1936e-02, -1.5046e-02],\n",
       "            [-1.8590e-02, -1.7201e-02, -1.9220e-02]],\n",
       " \n",
       "           [[-1.1719e-02, -1.1926e-02, -2.6681e-02],\n",
       "            [-2.1478e-02, -1.4142e-02, -1.0476e-02],\n",
       "            [-1.9153e-02, -1.6146e-02, -1.7522e-02]],\n",
       " \n",
       "           [[-1.4120e-02, -9.6434e-04, -2.4493e-03],\n",
       "            [-4.5751e-03, -3.6346e-02, -8.3142e-03],\n",
       "            [-1.2581e-02, -1.5882e-02, -2.6431e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-2.2747e-03, -1.6949e-02, -7.0367e-03],\n",
       "            [-8.4598e-03, -1.2638e-02, -1.0992e-02],\n",
       "            [ 9.2689e-04, -1.0148e-02, -1.2932e-02]],\n",
       " \n",
       "           [[ 2.3135e-03, -5.5956e-03, -2.1418e-02],\n",
       "            [-2.9278e-02, -1.4428e-03, -6.8166e-03],\n",
       "            [-7.6881e-03, -1.2255e-02, -1.0264e-02]],\n",
       " \n",
       "           [[-2.0622e-03, -1.1293e-03, -4.5414e-03],\n",
       "            [-3.9422e-03, -1.0161e-02, -2.5696e-03],\n",
       "            [-1.0729e-02, -4.1729e-03, -1.3555e-02]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[ 2.1522e-03,  9.3734e-03,  6.1744e-03],\n",
       "            [-1.9172e-02,  1.2538e-02,  2.0058e-04],\n",
       "            [ 1.3847e-02, -1.5014e-03, -4.2189e-03]],\n",
       " \n",
       "           [[ 3.9034e-03,  6.2036e-04,  9.9615e-03],\n",
       "            [ 3.6163e-03, -5.4916e-03,  3.8789e-03],\n",
       "            [ 1.7396e-02,  5.6892e-03,  1.1340e-02]],\n",
       " \n",
       "           [[-8.1911e-03, -2.1890e-02,  1.0171e-02],\n",
       "            [ 4.3713e-03,  2.2690e-03, -1.2439e-02],\n",
       "            [-1.6449e-02,  1.0848e-02,  5.5207e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 4.3041e-03,  3.1032e-03, -7.1255e-03],\n",
       "            [ 3.0167e-03,  3.8415e-03,  6.6839e-03],\n",
       "            [-8.9855e-03, -3.8018e-03, -1.0929e-02]],\n",
       " \n",
       "           [[-1.9557e-02,  1.8293e-02, -1.2693e-02],\n",
       "            [ 9.1237e-03, -7.0648e-03, -8.5851e-05],\n",
       "            [ 2.8720e-02,  2.8540e-02,  9.2053e-03]],\n",
       " \n",
       "           [[-2.7457e-03, -5.7138e-03,  1.6825e-03],\n",
       "            [-1.5145e-03,  2.4352e-04, -7.4389e-03],\n",
       "            [ 1.5225e-02,  3.7772e-03,  1.1642e-02]]],\n",
       " \n",
       " \n",
       "          [[[-1.0874e-02, -1.2541e-02,  2.2762e-02],\n",
       "            [-9.1565e-03, -4.2147e-03, -1.2612e-02],\n",
       "            [ 1.5452e-03, -5.5389e-03, -1.2146e-02]],\n",
       " \n",
       "           [[ 3.2202e-03, -1.1346e-02, -4.3072e-03],\n",
       "            [-4.1464e-03, -7.1322e-03, -7.1172e-04],\n",
       "            [-8.9004e-04, -1.5758e-03, -4.4856e-03]],\n",
       " \n",
       "           [[-1.4203e-02,  6.6294e-03, -4.8573e-03],\n",
       "            [-1.1029e-02, -1.3764e-02, -3.8058e-03],\n",
       "            [-4.7944e-03,  1.0183e-02,  9.3408e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.8399e-03,  4.3334e-03, -9.8100e-04],\n",
       "            [-8.2925e-03, -1.4949e-03,  6.3527e-03],\n",
       "            [-9.5980e-03, -1.6060e-03,  2.2435e-03]],\n",
       " \n",
       "           [[-5.3027e-03, -2.9961e-03, -1.9605e-02],\n",
       "            [ 1.0649e-02, -1.3858e-03, -6.4146e-03],\n",
       "            [ 6.6275e-03, -8.2005e-03, -5.4406e-03]],\n",
       " \n",
       "           [[-4.7704e-04, -7.0349e-03, -9.2217e-03],\n",
       "            [-5.2726e-03, -2.2603e-03, -1.1459e-03],\n",
       "            [ 2.6268e-03,  1.6528e-03,  2.9306e-04]]],\n",
       " \n",
       " \n",
       "          [[[-3.3415e-03, -9.0995e-03, -7.7234e-03],\n",
       "            [ 3.8506e-04, -1.4095e-03, -6.4783e-03],\n",
       "            [-9.0670e-03, -3.9695e-03, -1.5290e-02]],\n",
       " \n",
       "           [[-6.1763e-03, -4.8992e-03,  7.0965e-03],\n",
       "            [-6.0397e-03, -6.5023e-03, -1.0504e-02],\n",
       "            [ 7.0414e-03, -3.5348e-03, -2.4070e-03]],\n",
       " \n",
       "           [[-1.1901e-02, -3.4834e-03, -6.6724e-04],\n",
       "            [ 3.7544e-03, -1.2677e-02, -1.6825e-02],\n",
       "            [ 1.4285e-03,  4.8429e-03,  2.4788e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-1.8523e-02,  2.4462e-03,  7.9831e-03],\n",
       "            [ 6.3060e-03,  2.6712e-03, -2.2434e-03],\n",
       "            [-7.2605e-03, -6.1109e-03,  5.2013e-03]],\n",
       " \n",
       "           [[-2.7516e-04, -6.5982e-03, -1.5189e-03],\n",
       "            [-1.7955e-02, -1.2230e-02, -1.8675e-02],\n",
       "            [ 1.2905e-02, -8.9730e-04, -5.4949e-03]],\n",
       " \n",
       "           [[-4.6119e-03, -6.0973e-03,  1.7727e-03],\n",
       "            [-6.9878e-03, -7.0848e-03, -4.6150e-03],\n",
       "            [ 3.1063e-03,  1.0126e-02,  5.0778e-03]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-1.5258e-02, -5.8880e-03, -7.9041e-03],\n",
       "            [-1.6187e-02, -7.9725e-04, -1.0528e-02],\n",
       "            [ 4.1736e-04, -2.3769e-03, -3.4888e-03]],\n",
       " \n",
       "           [[-8.6106e-03, -1.1920e-02, -1.6545e-03],\n",
       "            [-5.4809e-03, -4.7607e-03, -2.1983e-03],\n",
       "            [ 6.1932e-04, -2.9722e-03, -5.8129e-03]],\n",
       " \n",
       "           [[-3.2493e-03, -1.5654e-02, -1.8734e-03],\n",
       "            [-3.0340e-03, -9.6709e-03, -1.5160e-02],\n",
       "            [-3.4641e-03,  4.7765e-03, -2.0464e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 3.2511e-05,  1.4401e-03,  5.2903e-03],\n",
       "            [ 1.0857e-03, -4.5238e-04, -3.2202e-03],\n",
       "            [ 1.9994e-03,  6.8081e-03, -3.5054e-03]],\n",
       " \n",
       "           [[-2.4603e-02, -3.5016e-04, -1.7366e-02],\n",
       "            [ 2.7711e-03, -1.2665e-02, -9.6932e-03],\n",
       "            [-4.0277e-03, -5.5065e-03, -1.2049e-02]],\n",
       " \n",
       "           [[-3.1217e-03, -6.5467e-03, -3.4776e-03],\n",
       "            [-1.2250e-03,  4.0852e-04, -6.9550e-03],\n",
       "            [-2.3849e-03, -1.3484e-03,  2.1084e-04]]],\n",
       " \n",
       " \n",
       "          [[[ 9.4275e-03, -2.3571e-03, -6.5518e-03],\n",
       "            [-3.1904e-03,  1.5644e-02, -9.2747e-04],\n",
       "            [-8.0833e-04,  6.6177e-03, -2.3581e-03]],\n",
       " \n",
       "           [[ 8.4713e-04, -1.2683e-02,  1.7981e-02],\n",
       "            [ 1.5963e-02, -8.6656e-03,  6.7392e-03],\n",
       "            [ 1.2325e-02, -2.5648e-03,  1.1410e-02]],\n",
       " \n",
       "           [[-7.4961e-03, -1.6450e-02,  4.4580e-03],\n",
       "            [ 2.7128e-02,  6.7999e-03, -7.2700e-03],\n",
       "            [-1.5589e-02,  2.3557e-02,  9.8973e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-6.6405e-03,  6.5041e-03, -5.2120e-03],\n",
       "            [-8.2190e-03,  1.1151e-02,  3.1185e-03],\n",
       "            [-9.8488e-03,  9.7755e-03, -1.3246e-03]],\n",
       " \n",
       "           [[ 3.4259e-03,  6.9614e-03, -8.3584e-04],\n",
       "            [ 8.7830e-03,  8.7905e-03, -3.4936e-03],\n",
       "            [ 1.1282e-02,  1.0216e-02,  3.1572e-03]],\n",
       " \n",
       "           [[ 3.9734e-03, -1.4567e-02,  4.4296e-03],\n",
       "            [ 1.3462e-02,  2.9456e-03, -1.2279e-03],\n",
       "            [ 3.4864e-03,  7.0279e-03,  1.0183e-02]]],\n",
       " \n",
       " \n",
       "          [[[ 2.7529e-03,  7.6572e-03, -2.8835e-03],\n",
       "            [ 3.2205e-04,  1.2858e-02,  1.7093e-03],\n",
       "            [ 8.5791e-03, -2.7925e-03,  1.0907e-02]],\n",
       " \n",
       "           [[ 1.7417e-03, -3.4122e-03,  1.1204e-02],\n",
       "            [ 1.0865e-02, -3.3947e-03,  8.6907e-03],\n",
       "            [ 3.5612e-04, -1.3693e-03,  4.1196e-03]],\n",
       " \n",
       "           [[ 1.6018e-02,  1.8038e-03, -5.3644e-05],\n",
       "            [ 5.9474e-03,  1.4150e-02, -2.8333e-03],\n",
       "            [-5.5335e-03,  2.0898e-02, -8.2229e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.9498e-03,  1.5480e-02, -3.2057e-03],\n",
       "            [-6.0396e-03,  3.5141e-03, -6.2109e-04],\n",
       "            [-6.8284e-03, -1.5612e-03,  2.7731e-03]],\n",
       " \n",
       "           [[-7.0876e-04,  5.9702e-03,  1.2527e-02],\n",
       "            [ 1.2470e-02,  5.6178e-03,  1.2912e-02],\n",
       "            [ 1.7086e-02, -8.0744e-03,  4.5485e-03]],\n",
       " \n",
       "           [[ 6.0994e-03, -2.3004e-03,  3.3459e-03],\n",
       "            [ 4.1729e-03,  2.5024e-03,  8.8080e-04],\n",
       "            [ 9.5588e-04, -3.8045e-03,  2.0317e-04]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[-1.1854e-02, -1.9760e-03,  1.6587e-03],\n",
       "            [-1.0577e-02, -5.1196e-03, -9.3856e-03],\n",
       "            [ 5.8460e-03, -7.8996e-03,  1.0958e-02]],\n",
       " \n",
       "           [[-3.7951e-03, -1.4808e-02, -3.9889e-03],\n",
       "            [-1.3036e-02, -2.4648e-03, -4.6519e-03],\n",
       "            [-9.2398e-04, -5.7262e-03,  3.8086e-03]],\n",
       " \n",
       "           [[-7.2548e-03,  6.6511e-03, -4.1026e-03],\n",
       "            [-7.7759e-03, -2.5763e-03,  8.6556e-03],\n",
       "            [-1.2531e-02, -7.1239e-03, -1.2429e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-6.7134e-03,  1.9077e-03, -9.2118e-03],\n",
       "            [-4.8671e-03, -7.7825e-04, -3.5164e-03],\n",
       "            [-4.3342e-03,  7.9811e-03, -2.0492e-03]],\n",
       " \n",
       "           [[ 9.5955e-03, -1.2196e-02, -5.2938e-03],\n",
       "            [ 5.9551e-03, -1.0344e-02,  8.9459e-03],\n",
       "            [-2.0473e-02,  1.6909e-02, -9.2412e-03]],\n",
       " \n",
       "           [[-1.7294e-03, -4.4400e-03, -7.6113e-04],\n",
       "            [ 1.7611e-03, -1.9876e-03, -2.4604e-03],\n",
       "            [-3.5384e-03, -1.0235e-02,  6.8409e-03]]],\n",
       " \n",
       " \n",
       "          [[[-9.6651e-03,  1.4164e-02, -2.2523e-02],\n",
       "            [ 2.9253e-03,  8.0465e-03, -1.0057e-03],\n",
       "            [-1.7496e-02, -2.5439e-03,  4.8847e-03]],\n",
       " \n",
       "           [[-1.0274e-03, -1.2741e-02,  2.6485e-04],\n",
       "            [ 1.0050e-02, -8.7205e-03,  7.9512e-03],\n",
       "            [-4.1085e-03,  2.0567e-03, -3.1933e-03]],\n",
       " \n",
       "           [[-1.6519e-03, -2.3419e-02,  1.0980e-02],\n",
       "            [ 5.3733e-03, -5.1568e-03, -1.2274e-03],\n",
       "            [-2.7404e-02,  1.5943e-02, -9.5904e-04]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.6727e-03, -9.6418e-04,  1.3769e-02],\n",
       "            [-1.2250e-02,  2.3159e-03, -1.4091e-04],\n",
       "            [ 1.0245e-02, -2.9058e-03, -3.6390e-03]],\n",
       " \n",
       "           [[ 8.4894e-03, -1.1009e-02, -1.3766e-02],\n",
       "            [-4.7911e-03,  7.0790e-03,  6.6413e-03],\n",
       "            [ 1.2627e-03, -8.3049e-03, -1.3258e-03]],\n",
       " \n",
       "           [[-5.6456e-03, -5.9690e-03, -4.5818e-03],\n",
       "            [ 8.2043e-03, -2.7991e-03,  4.0182e-03],\n",
       "            [-9.5465e-03,  6.5090e-03, -1.1300e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 8.1338e-04, -7.4302e-03,  1.5664e-02],\n",
       "            [ 2.5108e-03,  1.1384e-02,  3.4817e-03],\n",
       "            [ 4.7250e-03, -2.2938e-02, -2.3265e-04]],\n",
       " \n",
       "           [[-9.8959e-04, -1.7689e-03,  4.8457e-03],\n",
       "            [ 7.0810e-03,  1.2466e-02,  1.4046e-02],\n",
       "            [-6.9573e-03,  5.2206e-03,  1.8436e-03]],\n",
       " \n",
       "           [[ 7.6551e-03,  9.6048e-03,  3.5168e-04],\n",
       "            [ 1.3018e-02, -1.0880e-03, -2.8421e-03],\n",
       "            [-1.8866e-02,  1.6610e-02,  1.2309e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 5.5814e-03, -7.2157e-04,  1.0023e-03],\n",
       "            [-1.6296e-04,  1.0809e-03, -1.2501e-03],\n",
       "            [-1.0399e-03,  7.3776e-03, -3.4712e-03]],\n",
       " \n",
       "           [[ 4.5425e-03, -4.6767e-03, -4.0749e-04],\n",
       "            [ 4.3681e-03, -6.8178e-03,  1.7811e-02],\n",
       "            [ 4.0905e-03, -5.0280e-03,  3.3839e-03]],\n",
       " \n",
       "           [[-1.5401e-03, -8.1162e-03,  3.5470e-03],\n",
       "            [ 8.5088e-03,  7.8228e-03,  8.8447e-03],\n",
       "            [-2.8793e-03, -6.2154e-04,  5.3473e-03]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[ 1.2932e-03, -4.6813e-03, -1.3701e-03],\n",
       "            [-6.4253e-03,  1.2300e-02, -1.3539e-02],\n",
       "            [-3.7154e-03,  3.5443e-03,  1.1302e-02]],\n",
       " \n",
       "           [[-4.7883e-03, -6.0590e-03,  7.6164e-03],\n",
       "            [-9.8091e-03, -6.0090e-03,  1.7884e-02],\n",
       "            [ 1.4688e-02, -7.6592e-03,  9.3656e-03]],\n",
       " \n",
       "           [[-2.8023e-02,  7.0754e-03,  1.6617e-02],\n",
       "            [ 2.8633e-03, -5.3194e-03,  1.0285e-03],\n",
       "            [ 8.2911e-03,  1.1260e-03, -6.1192e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-3.9594e-03,  3.0334e-03, -1.1013e-02],\n",
       "            [ 1.2211e-03,  6.1552e-04,  3.7009e-04],\n",
       "            [ 1.2300e-03,  5.2052e-03, -1.5989e-02]],\n",
       " \n",
       "           [[-5.2952e-04,  1.2365e-02, -1.0176e-02],\n",
       "            [ 2.1364e-04, -2.1171e-02,  7.7597e-03],\n",
       "            [-1.4719e-02,  7.5527e-03,  3.5632e-03]],\n",
       " \n",
       "           [[-3.7664e-03, -7.0141e-04,  8.3710e-03],\n",
       "            [ 1.5058e-03, -3.8925e-03,  3.3963e-03],\n",
       "            [ 8.4995e-03, -6.5343e-03,  8.4658e-04]]],\n",
       " \n",
       " \n",
       "          [[[ 6.0925e-03,  2.2821e-03, -3.4602e-03],\n",
       "            [-5.2523e-04, -1.6435e-04, -4.5946e-03],\n",
       "            [ 9.9479e-04,  6.1363e-04, -1.0653e-02]],\n",
       " \n",
       "           [[-8.6203e-03,  2.7345e-03,  6.0422e-03],\n",
       "            [ 9.8013e-03, -4.9846e-03, -2.4178e-03],\n",
       "            [ 1.3672e-02, -7.3682e-03,  6.4676e-03]],\n",
       " \n",
       "           [[ 9.6187e-03, -2.4418e-02,  8.6640e-03],\n",
       "            [ 2.2732e-03, -1.7612e-03, -9.8052e-03],\n",
       "            [-1.1054e-02,  1.0758e-02,  7.4766e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[-4.9640e-03, -7.3309e-04,  9.0026e-03],\n",
       "            [-3.7261e-03,  9.7036e-04,  1.0253e-02],\n",
       "            [-1.4076e-02, -3.2454e-03,  5.2383e-03]],\n",
       " \n",
       "           [[ 3.7303e-03,  8.4426e-03, -1.4036e-02],\n",
       "            [ 3.3103e-03,  8.0030e-03, -2.0005e-02],\n",
       "            [ 2.5727e-02,  1.7897e-02,  6.3354e-03]],\n",
       " \n",
       "           [[-1.0617e-04, -7.0335e-03, -4.7284e-03],\n",
       "            [ 5.0238e-03,  1.5708e-03,  1.2912e-03],\n",
       "            [ 1.6763e-03,  2.2060e-03,  1.7586e-03]]],\n",
       " \n",
       " \n",
       "          [[[-6.1675e-03, -6.2885e-03, -1.5793e-02],\n",
       "            [-9.2106e-03, -5.0502e-03, -9.6853e-04],\n",
       "            [-2.9136e-02, -1.3752e-03, -8.5166e-03]],\n",
       " \n",
       "           [[-1.1076e-02, -6.1942e-03, -7.1282e-03],\n",
       "            [-7.3996e-03, -1.2729e-02, -5.9779e-03],\n",
       "            [-2.0253e-03, -6.4320e-03,  1.4161e-03]],\n",
       " \n",
       "           [[-1.8317e-02, -7.9753e-03, -5.9702e-03],\n",
       "            [ 5.4392e-03, -7.8321e-03, -3.1334e-03],\n",
       "            [-5.7128e-03, -1.1625e-02,  1.3545e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 2.9990e-03, -4.3636e-03, -5.7262e-03],\n",
       "            [-2.0839e-03, -2.4032e-03, -8.4257e-03],\n",
       "            [-1.0024e-02, -1.8154e-03, -1.3199e-03]],\n",
       " \n",
       "           [[-1.2905e-02, -1.2378e-02,  3.3683e-03],\n",
       "            [-2.5399e-02, -4.3312e-04, -9.1271e-03],\n",
       "            [ 5.5304e-04, -1.4312e-03,  6.0072e-04]],\n",
       " \n",
       "           [[-6.3098e-03, -1.8404e-03,  8.1911e-04],\n",
       "            [ 2.9934e-03, -7.3183e-03, -4.1527e-03],\n",
       "            [-9.8760e-03,  2.6361e-03,  4.0512e-03]]],\n",
       " \n",
       " \n",
       "          ...,\n",
       " \n",
       " \n",
       "          [[[ 2.9860e-03, -1.3428e-02,  6.0389e-03],\n",
       "            [-1.6377e-02, -7.1436e-04, -1.7421e-03],\n",
       "            [-6.1491e-03, -8.8047e-03, -1.6358e-03]],\n",
       " \n",
       "           [[ 3.0664e-03,  4.8386e-03, -1.2211e-02],\n",
       "            [-3.9862e-03, -5.8226e-03, -1.3207e-02],\n",
       "            [-2.3151e-03, -9.1576e-03, -1.1201e-02]],\n",
       " \n",
       "           [[-5.0776e-03,  4.7680e-03,  7.8976e-03],\n",
       "            [ 3.4710e-03, -1.5719e-02, -1.8716e-03],\n",
       "            [ 2.2722e-03, -1.4184e-04, -6.8929e-03]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 3.1818e-03, -4.2059e-04,  3.1528e-03],\n",
       "            [ 4.8684e-03, -7.2483e-03,  6.1024e-03],\n",
       "            [-7.0952e-03, -3.8955e-03,  4.4642e-03]],\n",
       " \n",
       "           [[-7.1924e-03,  2.1686e-02, -2.7262e-03],\n",
       "            [-6.7769e-03, -1.0232e-02, -2.6548e-03],\n",
       "            [-4.2881e-03,  7.0171e-03, -2.3289e-02]],\n",
       " \n",
       "           [[-3.9018e-03,  2.5764e-04,  1.4813e-03],\n",
       "            [-2.5282e-03, -6.0586e-03, -4.0777e-03],\n",
       "            [ 7.4701e-04, -5.5844e-03, -1.2242e-03]]],\n",
       " \n",
       " \n",
       "          [[[-7.7224e-03,  1.3067e-03, -8.5510e-03],\n",
       "            [ 3.8551e-02, -2.4588e-03,  1.5378e-02],\n",
       "            [ 6.3904e-03,  9.9196e-03,  7.0298e-03]],\n",
       " \n",
       "           [[-1.0988e-02,  5.6217e-03, -3.9511e-04],\n",
       "            [-2.5643e-03,  2.4081e-02,  1.1297e-02],\n",
       "            [ 6.8805e-03,  1.7608e-02,  1.4179e-02]],\n",
       " \n",
       "           [[-1.8857e-03, -2.2450e-05,  2.0657e-02],\n",
       "            [ 1.6083e-03,  4.6044e-03,  9.2683e-03],\n",
       "            [ 2.0676e-02,  1.1602e-02,  1.0017e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 4.7302e-03, -7.3856e-03,  1.4057e-02],\n",
       "            [ 4.7763e-03,  1.4419e-02, -1.7206e-03],\n",
       "            [ 6.8286e-03, -4.9877e-03,  2.3844e-03]],\n",
       " \n",
       "           [[ 2.0006e-02, -4.3308e-03, -1.1940e-02],\n",
       "            [ 8.0498e-03,  1.9384e-02,  1.4358e-02],\n",
       "            [ 3.4473e-03,  2.2159e-02, -9.2916e-03]],\n",
       " \n",
       "           [[-7.9338e-03, -4.2894e-05, -6.5351e-03],\n",
       "            [ 5.9668e-03,  5.0371e-03,  1.0700e-02],\n",
       "            [ 2.5707e-03,  7.4747e-03,  2.2371e-03]]],\n",
       " \n",
       " \n",
       "          [[[ 1.5250e-02,  4.6598e-03,  2.9180e-03],\n",
       "            [-2.7618e-02, -7.2547e-04, -1.6995e-02],\n",
       "            [-1.9267e-02,  1.7966e-02, -3.4086e-03]],\n",
       " \n",
       "           [[-9.0326e-03,  1.2470e-02,  1.0364e-03],\n",
       "            [ 9.9264e-03, -1.5882e-02, -3.9354e-03],\n",
       "            [ 7.1720e-03, -1.9307e-02,  5.9077e-04]],\n",
       " \n",
       "           [[ 6.4740e-03, -2.2429e-02,  2.6063e-03],\n",
       "            [ 6.9954e-03, -1.5507e-02,  7.8688e-03],\n",
       "            [ 1.5198e-02, -4.2601e-03, -2.1250e-02]],\n",
       " \n",
       "           ...,\n",
       " \n",
       "           [[ 4.5408e-03, -8.9228e-04, -3.8713e-03],\n",
       "            [-1.4822e-03, -8.6104e-03, -8.8642e-06],\n",
       "            [ 2.3870e-03,  7.9100e-03,  8.2933e-03]],\n",
       " \n",
       "           [[-2.2469e-02,  3.2958e-03,  9.1076e-03],\n",
       "            [-7.4517e-03,  1.3287e-02, -2.1030e-02],\n",
       "            [-9.9789e-03, -8.4374e-03,  1.0511e-02]],\n",
       " \n",
       "           [[-3.6823e-03,  5.6951e-04,  5.9617e-04],\n",
       "            [ 2.5015e-03, -5.0385e-03, -1.4640e-03],\n",
       "            [-5.9543e-04, -4.0804e-03, -6.8835e-03]]]]], device='cuda:0'),\n",
       " 'conv2.bias': tensor([[-0.0658,  0.0079, -0.0329,  ...,  0.0159, -0.0105, -0.0082],\n",
       "         [ 0.0364,  0.0086,  0.0249,  ..., -0.0041, -0.0123,  0.0061],\n",
       "         [ 0.0142,  0.0196,  0.0226,  ..., -0.0121, -0.0076,  0.0205],\n",
       "         ...,\n",
       "         [-0.0239,  0.0312, -0.0451,  ...,  0.0184, -0.0123, -0.0147],\n",
       "         [-0.0155,  0.0083,  0.0142,  ..., -0.0085,  0.0017,  0.0084],\n",
       "         [-0.0005,  0.0039, -0.0157,  ..., -0.0128,  0.0271, -0.0044]],\n",
       "        device='cuda:0'),\n",
       " 'fc1.weight': tensor([[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0067,  0.0018,  ...,  0.0031,  0.0049,  0.0066],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [ 0.0000,  0.0191,  0.0051,  ...,  0.0090,  0.0141,  0.0189],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0000, -0.0071, -0.0019,  ..., -0.0034, -0.0053, -0.0071]],\n",
       " \n",
       "         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0116, -0.0082, -0.0029,  ..., -0.0155, -0.0119, -0.0116],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [ 0.0117,  0.0082,  0.0029,  ...,  0.0156,  0.0120,  0.0116],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0156, -0.0110, -0.0038,  ..., -0.0209, -0.0160, -0.0155]],\n",
       " \n",
       "         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0020, -0.0009, -0.0023,  ..., -0.0045, -0.0027, -0.0020],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [ 0.0071,  0.0031,  0.0081,  ...,  0.0160,  0.0098,  0.0073],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],\n",
       " \n",
       "         ...,\n",
       " \n",
       "         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [ 0.0156,  0.0000,  0.0132,  ...,  0.0286,  0.0244,  0.0265],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],\n",
       " \n",
       "         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0015, -0.0011, -0.0009,  ..., -0.0036, -0.0031, -0.0071],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [ 0.0060,  0.0046,  0.0036,  ...,  0.0150,  0.0130,  0.0295],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0083, -0.0064, -0.0050,  ..., -0.0208, -0.0180, -0.0409]],\n",
       " \n",
       "         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [-0.0372, -0.0280, -0.0124,  ..., -0.0406, -0.0224, -0.0101],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [-0.0189, -0.0142, -0.0063,  ..., -0.0205, -0.0114, -0.0051],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],\n",
       "        device='cuda:0'),\n",
       " 'fc1.bias': tensor([[ 0.0000,  0.0185,  0.0000,  ...,  0.0529,  0.0000, -0.0198],\n",
       "         [ 0.0000, -0.0441,  0.0000,  ...,  0.0443,  0.0000, -0.0594],\n",
       "         [ 0.0000, -0.0122,  0.0000,  ...,  0.0438,  0.0000,  0.0000],\n",
       "         ...,\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0524,  0.0000,  0.0000],\n",
       "         [ 0.0000, -0.0106,  0.0000,  ...,  0.0440,  0.0000, -0.0610],\n",
       "         [ 0.0000, -0.0933,  0.0000,  ..., -0.0473,  0.0000,  0.0000]],\n",
       "        device='cuda:0'),\n",
       " 'fc2.weight': tensor([[[ 0.0000e+00,  3.6354e-03,  0.0000e+00,  ...,  3.3698e-02,\n",
       "            0.0000e+00,  3.1801e-03],\n",
       "          [ 0.0000e+00,  3.5352e-03,  0.0000e+00,  ...,  3.2769e-02,\n",
       "            0.0000e+00,  3.0924e-03],\n",
       "          [ 0.0000e+00,  3.5016e-03,  0.0000e+00,  ...,  3.2458e-02,\n",
       "            0.0000e+00,  3.0630e-03],\n",
       "          ...,\n",
       "          [-0.0000e+00, -3.1660e-02, -0.0000e+00,  ..., -2.9346e-01,\n",
       "           -0.0000e+00, -2.7694e-02],\n",
       "          [ 0.0000e+00,  3.1316e-03,  0.0000e+00,  ...,  2.9028e-02,\n",
       "            0.0000e+00,  2.7393e-03],\n",
       "          [ 0.0000e+00,  4.2018e-03,  0.0000e+00,  ...,  3.8948e-02,\n",
       "            0.0000e+00,  3.6755e-03]],\n",
       " \n",
       "         [[ 0.0000e+00,  7.2065e-03,  0.0000e+00,  ...,  2.9356e-02,\n",
       "            0.0000e+00,  8.7264e-03],\n",
       "          [ 0.0000e+00,  6.6121e-03,  0.0000e+00,  ...,  2.6935e-02,\n",
       "            0.0000e+00,  8.0067e-03],\n",
       "          [ 0.0000e+00,  6.7659e-03,  0.0000e+00,  ...,  2.7562e-02,\n",
       "            0.0000e+00,  8.1929e-03],\n",
       "          ...,\n",
       "          [ 0.0000e+00,  7.0268e-03,  0.0000e+00,  ...,  2.8624e-02,\n",
       "            0.0000e+00,  8.5088e-03],\n",
       "          [ 0.0000e+00,  5.7157e-03,  0.0000e+00,  ...,  2.3283e-02,\n",
       "            0.0000e+00,  6.9212e-03],\n",
       "          [ 0.0000e+00,  8.1111e-03,  0.0000e+00,  ...,  3.3041e-02,\n",
       "            0.0000e+00,  9.8218e-03]],\n",
       " \n",
       "         [[ 0.0000e+00,  4.3953e-03,  0.0000e+00,  ...,  3.9794e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  3.8886e-03,  0.0000e+00,  ...,  3.5206e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  3.9954e-03,  0.0000e+00,  ...,  3.6174e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          ...,\n",
       "          [ 0.0000e+00,  4.2279e-03,  0.0000e+00,  ...,  3.8279e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [-0.0000e+00, -3.6841e-02, -0.0000e+00,  ..., -3.3356e-01,\n",
       "           -0.0000e+00, -0.0000e+00],\n",
       "          [ 0.0000e+00,  4.7483e-03,  0.0000e+00,  ...,  4.2990e-02,\n",
       "            0.0000e+00,  0.0000e+00]],\n",
       " \n",
       "         ...,\n",
       " \n",
       "         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  3.9412e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  3.6385e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  3.6365e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          ...,\n",
       "          [-0.0000e+00, -0.0000e+00, -0.0000e+00,  ..., -3.3984e-01,\n",
       "           -0.0000e+00, -0.0000e+00],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  3.3791e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  4.4862e-02,\n",
       "            0.0000e+00,  0.0000e+00]],\n",
       " \n",
       "         [[ 0.0000e+00,  1.9226e-02,  0.0000e+00,  ...,  2.3166e-02,\n",
       "            0.0000e+00,  2.1841e-04],\n",
       "          [ 0.0000e+00,  1.9772e-02,  0.0000e+00,  ...,  2.3824e-02,\n",
       "            0.0000e+00,  2.2461e-04],\n",
       "          [ 0.0000e+00,  1.8173e-02,  0.0000e+00,  ...,  2.1897e-02,\n",
       "            0.0000e+00,  2.0644e-04],\n",
       "          ...,\n",
       "          [ 0.0000e+00,  1.9233e-02,  0.0000e+00,  ...,  2.3175e-02,\n",
       "            0.0000e+00,  2.1849e-04],\n",
       "          [-0.0000e+00, -1.7064e-01, -0.0000e+00,  ..., -2.0560e-01,\n",
       "           -0.0000e+00, -1.9385e-03],\n",
       "          [ 0.0000e+00,  2.3099e-02,  0.0000e+00,  ...,  2.7833e-02,\n",
       "            0.0000e+00,  2.6241e-04]],\n",
       " \n",
       "         [[ 0.0000e+00,  8.4050e-03,  0.0000e+00,  ...,  3.4529e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [-0.0000e+00, -7.3317e-02, -0.0000e+00,  ..., -3.0120e-01,\n",
       "           -0.0000e+00, -0.0000e+00],\n",
       "          [ 0.0000e+00,  7.7367e-03,  0.0000e+00,  ...,  3.1784e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          ...,\n",
       "          [ 0.0000e+00,  8.5000e-03,  0.0000e+00,  ...,  3.4920e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  7.2293e-03,  0.0000e+00,  ...,  2.9699e-02,\n",
       "            0.0000e+00,  0.0000e+00],\n",
       "          [ 0.0000e+00,  9.7435e-03,  0.0000e+00,  ...,  4.0028e-02,\n",
       "            0.0000e+00,  0.0000e+00]]], device='cuda:0'),\n",
       " 'fc2.bias': tensor([[ 0.1029,  0.1001,  0.0991,  0.1055,  0.0894,  0.0929,  0.0989, -0.8964,\n",
       "           0.0887,  0.1190],\n",
       "         [ 0.1074,  0.0986,  0.1009, -0.8951,  0.0883,  0.0894,  0.0998,  0.1047,\n",
       "           0.0852,  0.1209],\n",
       "         [ 0.1090,  0.0965,  0.0991,  0.1052,  0.0889,  0.0926,  0.1000,  0.1049,\n",
       "          -0.9138,  0.1178],\n",
       "         [ 0.1027,  0.1040,  0.0957,  0.1063,  0.0901, -0.9088,  0.1001,  0.1013,\n",
       "           0.0894,  0.1190],\n",
       "         [ 0.1010,  0.1038,  0.0960, -0.8971,  0.0902,  0.0931,  0.1011,  0.1039,\n",
       "           0.0905,  0.1175],\n",
       "         [ 0.1032,  0.1013,  0.0967,  0.1100,  0.0943,  0.0927,  0.0957, -0.9006,\n",
       "           0.0872,  0.1194],\n",
       "         [ 0.1029,  0.1021,  0.0985,  0.1073,  0.0816,  0.0899,  0.1027,  0.1034,\n",
       "          -0.9092,  0.1208],\n",
       "         [ 0.1050,  0.1017,  0.0949,  0.1064,  0.0921,  0.0923,  0.0977,  0.1038,\n",
       "           0.0831, -0.8768],\n",
       "         [ 0.1112,  0.1046,  0.0985,  0.1082,  0.0886,  0.0904, -0.9047,  0.0976,\n",
       "           0.0862,  0.1194],\n",
       "         [ 0.1021,  0.0996,  0.1009,  0.1108,  0.0926,  0.0911, -0.9052,  0.1033,\n",
       "           0.0848,  0.1199],\n",
       "         [ 0.1020,  0.0982,  0.0988,  0.1120,  0.0856,  0.0944,  0.1000,  0.1002,\n",
       "           0.0875, -0.8788],\n",
       "         [ 0.1041,  0.1004, -0.9045,  0.1082,  0.0850,  0.0956,  0.1010,  0.1058,\n",
       "           0.0855,  0.1189],\n",
       "         [ 0.1045,  0.1024,  0.0977,  0.1112,  0.0871, -0.9073,  0.0988,  0.1014,\n",
       "           0.0861,  0.1182],\n",
       "         [ 0.1057,  0.0977,  0.0954,  0.1040,  0.0904,  0.0929,  0.0996, -0.8965,\n",
       "           0.0878,  0.1232],\n",
       "         [ 0.1053,  0.1031,  0.0960,  0.1070, -0.9091,  0.0926,  0.0995,  0.1057,\n",
       "           0.0854,  0.1145],\n",
       "         [-0.9011,  0.0973,  0.1044,  0.1075,  0.0877,  0.0899,  0.1021,  0.1039,\n",
       "           0.0879,  0.1203],\n",
       "         [ 0.0977, -0.9023,  0.0950,  0.1147,  0.0903,  0.0922,  0.0999,  0.1046,\n",
       "           0.0904,  0.1175],\n",
       "         [ 0.1038,  0.1046, -0.9020,  0.1059,  0.0906,  0.0902,  0.0985,  0.0996,\n",
       "           0.0885,  0.1201],\n",
       "         [ 0.1055,  0.1013,  0.0973,  0.1063,  0.0920,  0.0911, -0.9001,  0.1029,\n",
       "           0.0843,  0.1193],\n",
       "         [ 0.1049, -0.8978,  0.0952,  0.1068,  0.0913,  0.0886,  0.0982,  0.1000,\n",
       "           0.0931,  0.1198],\n",
       "         [ 0.1073,  0.1017,  0.0947,  0.1098,  0.0871,  0.0919,  0.1010,  0.1003,\n",
       "          -0.9138,  0.1199],\n",
       "         [ 0.0986,  0.1003,  0.0999,  0.1085, -0.9126,  0.0928,  0.1020,  0.1034,\n",
       "           0.0867,  0.1204],\n",
       "         [ 0.1024,  0.1021,  0.0975,  0.1063,  0.0905, -0.9066,  0.1014,  0.1021,\n",
       "           0.0862,  0.1181],\n",
       "         [ 0.1031,  0.1012,  0.0915,  0.1092,  0.0918,  0.0923,  0.0990,  0.1062,\n",
       "           0.0864, -0.8807],\n",
       "         [ 0.1022,  0.1030,  0.0968,  0.1103,  0.0878,  0.0888,  0.0999,  0.1038,\n",
       "           0.0866, -0.8794],\n",
       "         [ 0.1048,  0.1047,  0.0974,  0.1098,  0.0897,  0.0868, -0.9000,  0.1011,\n",
       "           0.0840,  0.1217],\n",
       "         [ 0.1027,  0.1019,  0.0970,  0.1098,  0.0879,  0.0917,  0.0980,  0.1001,\n",
       "           0.0883, -0.8773],\n",
       "         [ 0.1040,  0.1006,  0.0940,  0.1093,  0.0888, -0.9069,  0.0999,  0.1020,\n",
       "           0.0849,  0.1234],\n",
       "         [ 0.1004,  0.0999,  0.1017, -0.8947,  0.0916,  0.0929,  0.0978,  0.1087,\n",
       "           0.0858,  0.1158],\n",
       "         [ 0.1040,  0.1027,  0.0981,  0.1104,  0.0914,  0.0889, -0.9054,  0.1016,\n",
       "           0.0847,  0.1236],\n",
       "         [ 0.1020,  0.0998,  0.0962,  0.1072,  0.0868,  0.0934,  0.1044,  0.1043,\n",
       "           0.0897, -0.8837],\n",
       "         [ 0.1027,  0.1063,  0.0949,  0.1100, -0.9101,  0.0905,  0.0957,  0.1008,\n",
       "           0.0855,  0.1237],\n",
       "         [ 0.1040,  0.1025,  0.0966,  0.1088,  0.0865,  0.0917,  0.1012,  0.0990,\n",
       "          -0.9094,  0.1191],\n",
       "         [ 0.1035,  0.1053,  0.0993,  0.1088,  0.0879,  0.0924,  0.0944, -0.8966,\n",
       "           0.0830,  0.1219],\n",
       "         [ 0.1044,  0.1002,  0.0987,  0.1056,  0.0894,  0.0914,  0.1022,  0.0996,\n",
       "          -0.9136,  0.1222],\n",
       "         [ 0.1025,  0.1017,  0.0945,  0.1074,  0.0848,  0.0922,  0.1037,  0.1063,\n",
       "           0.0863, -0.8794],\n",
       "         [ 0.1024,  0.0992,  0.1013, -0.8946,  0.0892,  0.0951,  0.0964,  0.1066,\n",
       "           0.0861,  0.1183],\n",
       "         [ 0.1040,  0.1037,  0.0939,  0.1088, -0.9115,  0.0930,  0.0978,  0.1027,\n",
       "           0.0861,  0.1215],\n",
       "         [ 0.0993,  0.1001, -0.9085,  0.1098,  0.0904,  0.0946,  0.1037,  0.1017,\n",
       "           0.0888,  0.1201],\n",
       "         [ 0.1040,  0.1011,  0.0975,  0.1055,  0.0864,  0.0968,  0.0988,  0.1022,\n",
       "          -0.9149,  0.1226],\n",
       "         [ 0.1043,  0.1017,  0.0961, -0.8940,  0.0899,  0.0934,  0.1010,  0.1042,\n",
       "           0.0859,  0.1175],\n",
       "         [ 0.1007, -0.9011,  0.0980,  0.1102,  0.0853,  0.0916,  0.1019,  0.1046,\n",
       "           0.0886,  0.1204],\n",
       "         [ 0.1041,  0.0981,  0.1004,  0.1073,  0.0867,  0.0928,  0.1012,  0.1011,\n",
       "          -0.9139,  0.1222],\n",
       "         [ 0.1034,  0.0961,  0.0999, -0.8931,  0.0906,  0.0943,  0.0992,  0.1050,\n",
       "           0.0869,  0.1176],\n",
       "         [ 0.0994,  0.1022,  0.1007,  0.1087,  0.0900,  0.0935,  0.1012,  0.1025,\n",
       "           0.0872, -0.8853],\n",
       "         [ 0.1058,  0.1006,  0.0960,  0.1072,  0.0864,  0.0927,  0.1016,  0.1051,\n",
       "           0.0877, -0.8830],\n",
       "         [ 0.1056,  0.1002,  0.0980,  0.1072,  0.0891,  0.0938,  0.1003, -0.8990,\n",
       "           0.0874,  0.1175],\n",
       "         [ 0.1031,  0.1026,  0.0952,  0.1120, -0.9075,  0.0902,  0.0956,  0.1025,\n",
       "           0.0855,  0.1208],\n",
       "         [ 0.1040,  0.1032,  0.0983,  0.1097,  0.0930,  0.0968,  0.0954,  0.1001,\n",
       "          -0.9137,  0.1132],\n",
       "         [ 0.1009,  0.1003,  0.0978,  0.1071,  0.0893, -0.9081,  0.0999,  0.1053,\n",
       "           0.0866,  0.1209],\n",
       "         [ 0.1007,  0.0994,  0.0993,  0.1102,  0.0886,  0.0946, -0.9005,  0.1009,\n",
       "           0.0879,  0.1189],\n",
       "         [ 0.1006, -0.8974,  0.0991,  0.1077,  0.0871,  0.0901,  0.0957,  0.1055,\n",
       "           0.0860,  0.1257],\n",
       "         [ 0.1054,  0.1029,  0.0957,  0.1102,  0.0856,  0.0943,  0.0970, -0.8972,\n",
       "           0.0849,  0.1212],\n",
       "         [ 0.1010,  0.0971,  0.1018,  0.1086,  0.0868,  0.0910,  0.1012,  0.1038,\n",
       "          -0.9112,  0.1200],\n",
       "         [ 0.1023,  0.0999,  0.0982,  0.1091,  0.0834,  0.0909, -0.8964,  0.1058,\n",
       "           0.0888,  0.1181],\n",
       "         [-0.8968,  0.1004,  0.1047,  0.1084,  0.0892,  0.0936,  0.1005,  0.1010,\n",
       "           0.0834,  0.1155],\n",
       "         [ 0.1009,  0.0982, -0.8969,  0.1056,  0.0875,  0.0955,  0.1002,  0.1062,\n",
       "           0.0828,  0.1200],\n",
       "         [ 0.1084,  0.1022,  0.0965,  0.1050,  0.0921,  0.0918,  0.0983,  0.1048,\n",
       "           0.0846, -0.8838],\n",
       "         [-0.9004,  0.1021,  0.0921,  0.1083,  0.0874,  0.0901,  0.1018,  0.1054,\n",
       "           0.0883,  0.1249],\n",
       "         [ 0.1058,  0.1031,  0.0970, -0.8942,  0.0911,  0.0902,  0.1015,  0.1003,\n",
       "           0.0881,  0.1171],\n",
       "         [ 0.1041,  0.0988,  0.0981, -0.8925,  0.0887,  0.0928,  0.0990,  0.1045,\n",
       "           0.0862,  0.1203],\n",
       "         [ 0.1035,  0.0956,  0.0955,  0.1056,  0.0877,  0.0958,  0.1023, -0.8926,\n",
       "           0.0887,  0.1178],\n",
       "         [ 0.1029,  0.1058,  0.0972,  0.1060,  0.0871,  0.0875,  0.1000,  0.1029,\n",
       "          -0.9130,  0.1236],\n",
       "         [ 0.1029, -0.8975,  0.0947,  0.1036,  0.0904,  0.0920,  0.1021,  0.1040,\n",
       "           0.0885,  0.1193]], device='cuda:0')}"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# A function computing the gradient of compute_loss with respect to its first argument\n",
    "ft_compute_grad = grad(compute_loss)\n",
    "ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0, 0))\n",
    "ft_per_sample_grads = ft_compute_sample_grad(params, buffers, data, targets)\n",
    "ft_per_sample_grads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'])"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ft_per_sample_grads.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`in_dims`?\n",
    "- Denotes to PyTorch which dimensions to apply the function along\n",
    "    - Similar to indexing those dimensions and then passing the result into the function being vmap'd\n",
    "- None indicates no map dimension--i.e., don't map the function over these\n",
    "- `in_dim` should have a structur reflecting the inputs themselves, selecting dimensions\n",
    "- In this case:\n",
    "    - None - Don't iterate over params; treat them as a constant wrt summation indexj\n",
    "    - None - Don't iterate over buffers; they're a constant with respect to iteration\n",
    "    - 0 - The first dimension of data is the batch dimension\n",
    "    - 0 - ...which should be aligned with the first dimension of the parameters\n",
    "\n",
    "[Documentation about vmap](https://pytorch.org/docs/stable/generated/torch.vmap.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can be thought of as\n",
    "$$\\sum _ {i \\in \\Gamma_\\text{batch}} \\text{compute\\_loss}(\\text{params}, \\text{buffers}, \\text{samples}[i], \\text{targets}[i])$$\n",
    "\n",
    "The value of `in_dims` places the $i$ in the index-brackets of the input tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Performance of `vmap`\n",
    "- Much faster (documentation claims order of magnitude) than computing in a simple for loop\n",
    "- TODO: where is this speedup coming from?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What should the function that computes $\\frac{d\\ln P}{d\\theta}$ look like?\n",
    "\n",
    "- Does sampling need to be accounted for in the \"differentiation\"?\n",
    "- Argument: no! We have an energy gradient and we've performed importance sampling so that we're finding the gradient at important points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def d_ln_P_d_theta(self, psi_x: torch.Tensor)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tqs2",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}