1.TrainStudentMSE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a1a865bf-d19f-4803-b256-4a3247bd4667",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA: True cuda:0\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from time import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "\n",
    "\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from scipy.sparse.linalg import eigs\n",
    "\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "DEVICE = torch.device('cuda:0')\n",
    "print(\"CUDA:\", USE_CUDA, DEVICE)\n",
    "\n",
    "from tensorboardX import SummaryWriter\n",
    "sw = SummaryWriter(logdir='./Student_MSE', flush_secs=5)\n",
    "\n",
    "import math\n",
    "from typing import Optional, List, Union\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.typing import OptTensor\n",
    "from torch_geometric.nn.conv import MessagePassing\n",
    "from torch_geometric.transforms import LaplacianLambdaMax\n",
    "from torch_geometric.utils import remove_self_loops, add_self_loops, get_laplacian\n",
    "from torch_geometric.utils import to_dense_adj\n",
    "from torch_scatter import scatter_add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "53d26d62-aed2-4d96-bc93-9efda7479480",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_graphdata_channel1(graph_signal_matrix_filename, num_of_hours, num_of_days, num_of_weeks, batch_size,\n",
    "                            shuffle=True, DEVICE = torch.device('cuda:0')):\n",
    "    '''\n",
    "    :param graph_signal_matrix_filename: str\n",
    "    :param num_of_hours: int\n",
    "    :param num_of_days: int\n",
    "    :param num_of_weeks: int\n",
    "    :param DEVICE:\n",
    "    :param batch_size: int\n",
    "    :return:\n",
    "    three DataLoaders, each dataloader contains:\n",
    "    test_x_tensor: (B, N_nodes, in_feature, T_input)\n",
    "    test_decoder_input_tensor: (B, N_nodes, T_output)\n",
    "    test_target_tensor: (B, N_nodes, T_output)\n",
    "    '''\n",
    "\n",
    "    file = os.path.basename(graph_signal_matrix_filename).split('.')[0]\n",
    "    filename = os.path.join('./data/PEMS04/', file + '_r' + str(num_of_hours) + '_d' + str(num_of_days) + '_w' + str(num_of_weeks)) +'_astcgn'\n",
    "    print('load file:', filename)\n",
    "\n",
    "    file_data = np.load(filename + '.npz')\n",
    "    train_x = file_data['train_x']  # (10181, 307, 3, 12)\n",
    "    train_x = train_x[:, :, 0:1, :]\n",
    "    train_target = file_data['train_target']  # (10181, 307, 12)\n",
    "\n",
    "    val_x = file_data['val_x']\n",
    "    val_x = val_x[:, :, 0:1, :]\n",
    "    val_target = file_data['val_target']\n",
    "\n",
    "    test_x = file_data['test_x']\n",
    "    test_x = test_x[:, :, 0:1, :]\n",
    "    test_target = file_data['test_target']\n",
    "\n",
    "    mean = file_data['mean'][:, :, 0:1, :]  # (1, 1, 3, 1)\n",
    "    std = file_data['std'][:, :, 0:1, :]  # (1, 1, 3, 1)\n",
    "\n",
    "    # ------- train_loader -------\n",
    "    train_x_tensor = torch.from_numpy(train_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)\n",
    "    train_target_tensor = torch.from_numpy(train_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)\n",
    "    train_dataset = torch.utils.data.TensorDataset(train_x_tensor, train_target_tensor)\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle)\n",
    "\n",
    "    # ------- val_loader -------\n",
    "    val_x_tensor = torch.from_numpy(val_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)\n",
    "    val_target_tensor = torch.from_numpy(val_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)\n",
    "    val_dataset = torch.utils.data.TensorDataset(val_x_tensor, val_target_tensor)\n",
    "    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # ------- test_loader -------\n",
    "    test_x_tensor = torch.from_numpy(test_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)\n",
    "    test_target_tensor = torch.from_numpy(test_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)\n",
    "    test_dataset = torch.utils.data.TensorDataset(test_x_tensor, test_target_tensor)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # print\n",
    "    print('train:', train_x_tensor.size(), train_target_tensor.size())\n",
    "    print('val:', val_x_tensor.size(), val_target_tensor.size())\n",
    "    print('test:', test_x_tensor.size(), test_target_tensor.size())\n",
    "\n",
    "    return train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, mean, std\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3913859a-5cc7-4103-803c-c01c07cc624b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load file: ./data/PEMS04/pems04_r1_d1_w1_astcgn\n",
      "train: torch.Size([13575, 307, 1, 12]) torch.Size([13575, 307, 12])\n",
      "val: torch.Size([1697, 307, 1, 12]) torch.Size([1697, 307, 12])\n",
      "test: torch.Size([1697, 307, 1, 12]) torch.Size([1697, 307, 12])\n"
     ]
    }
   ],
   "source": [
    "graph_signal_matrix_filename = './data/PEMS04/pems04.npz'\n",
    "batch_size = 16\n",
    "num_of_weeks = 1\n",
    "num_of_days = 1\n",
    "num_of_hours = 1\n",
    "\n",
    "train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, _mean, _std = load_graphdata_channel1(\n",
    "    graph_signal_matrix_filename, num_of_hours, num_of_days, num_of_weeks, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eba75207-cb40-4e13-b240-b5db01af551c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_adjacency_matrix(distance_df_filename, num_of_vertices, id_filename=None):\n",
    "    '''\n",
    "    Parameters\n",
    "    ----------\n",
    "    distance_df_filename: str, path of the csv file contains edges information\n",
    "    num_of_vertices: int, the number of vertices\n",
    "    Returns\n",
    "    ----------\n",
    "    A: np.ndarray, adjacency matrix\n",
    "    '''\n",
    "    if 'npy' in distance_df_filename:  # false\n",
    "        adj_mx = np.load(distance_df_filename)\n",
    "        return adj_mx, None\n",
    "    else:\n",
    "        \n",
    "        #--------------------------------------------- read from here\n",
    "        import csv\n",
    "        A = np.zeros((int(num_of_vertices), int(num_of_vertices)),dtype=np.float32)\n",
    "        distaneA = np.zeros((int(num_of_vertices), int(num_of_vertices)), dtype=np.float32)\n",
    "\n",
    "        #------------ Ignore\n",
    "        if id_filename: # false\n",
    "            with open(id_filename, 'r') as f:\n",
    "                id_dict = {int(i): idx for idx, i in enumerate(f.read().strip().split('\\n'))}  # 把节点id（idx）映射成从0开始的索引\n",
    "\n",
    "            with open(distance_df_filename, 'r') as f:\n",
    "                f.readline()\n",
    "                reader = csv.reader(f)\n",
    "                for row in reader:\n",
    "                    if len(row) != 3:\n",
    "                        continue\n",
    "                    i, j, distance = int(row[0]), int(row[1]), float(row[2])\n",
    "                    A[id_dict[i], id_dict[j]] = 1\n",
    "                    distaneA[id_dict[i], id_dict[j]] = distance\n",
    "            return A, distaneA\n",
    "\n",
    "        else:\n",
    "         #-------------Continue reading\n",
    "            with open(distance_df_filename, 'r') as f:\n",
    "                f.readline()\n",
    "                reader = csv.reader(f)\n",
    "                for row in reader:\n",
    "                    if len(row) != 3:\n",
    "                        continue\n",
    "                    i, j, distance = int(row[0]), int(row[1]), float(row[2])\n",
    "                    A[i, j] = 1\n",
    "                    distaneA[i, j] = distance\n",
    "            return A, distaneA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c5a8d098-f629-4513-8dd3-a7a8def987b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_filename = None\n",
    "adj_filename = './data/PEMS04/PEMS04.csv'\n",
    "num_of_vertices = 307\n",
    "adj_mx, distance_mx = get_adjacency_matrix(adj_filename, num_of_vertices, id_filename) #  adj_mx and distance_mx (307, 307)\n",
    "rows, cols = np.where(adj_mx == 1)\n",
    "edges = zip(rows.tolist(), cols.tolist())\n",
    "gr = nx.Graph()\n",
    "gr.add_edges_from(edges)\n",
    "# nx.draw(gr, node_size=3)\n",
    "plt.show()\n",
    "rows, cols = np.where(adj_mx == 1)\n",
    "edges = zip(rows.tolist(), cols.tolist())\n",
    "edge_index_data = torch.LongTensor(np.array([rows, cols])).to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a9d96052-6506-427c-85fe-a62e35748109",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.ASTGCN import ASTGCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ebcc7a3f-54c0-4732-9320-a48a0a960f13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ASTGCN(\n",
      "  (_blocklist): ModuleList(\n",
      "    (0): ASTGCNBlock(\n",
      "      (_temporal_attention): TemporalAttention()\n",
      "      (_spatial_attention): SpatialAttention()\n",
      "      (_chebconv_attention): ChebConvAttention(1, 64, K=3, normalization=None)\n",
      "      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))\n",
      "      (_residual_convolution): Conv2d(1, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
      "      (_dropout): Dropout2d(p=0.01, inplace=True)\n",
      "    )\n",
      "    (1): ASTGCNBlock(\n",
      "      (_temporal_attention): TemporalAttention()\n",
      "      (_spatial_attention): SpatialAttention()\n",
      "      (_chebconv_attention): ChebConvAttention(64, 64, K=3, normalization=None)\n",
      "      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))\n",
      "      (_residual_convolution): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
      "      (_dropout): Dropout2d(p=0.01, inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (_final_conv): Conv2d(12, 12, kernel_size=(1, 64), stride=(1, 1))\n",
      ")\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "nb_block = 2\n",
    "in_channels = 1\n",
    "K = 3\n",
    "nb_chev_filter =64\n",
    "nb_time_filter = 64\n",
    "time_strides = num_of_hours\n",
    "num_for_predict = 12\n",
    "len_input = 12\n",
    "dropout=0.01\n",
    "\n",
    "#L_tilde = scaled_Laplacian(adj_mx)\n",
    "#cheb_polynomials = [torch.from_numpy(i).type(torch.FloatTensor).to(DEVICE) for i in cheb_polynomial(L_tilde, K)]\n",
    "# net = ASTGCN( nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, num_for_predict, len_input, num_of_vertices).to(DEVICE)\n",
    "Teacher=ASTGCN( nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, num_for_predict, len_input, num_of_vertices,dropout).to(DEVICE)\n",
    "print(Teacher)\n",
    "print(time_strides)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "50dee06c-5f15-457b-80bf-e9538b7792a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ASTGCN(\n",
      "  (_blocklist): ModuleList(\n",
      "    (0): ASTGCNBlock(\n",
      "      (_temporal_attention): TemporalAttention()\n",
      "      (_spatial_attention): SpatialAttention()\n",
      "      (_chebconv_attention): ChebConvAttention(1, 64, K=3, normalization=None)\n",
      "      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))\n",
      "      (_residual_convolution): Conv2d(1, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
      "      (_dropout): Dropout2d(p=0.01, inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (_final_conv): Conv2d(12, 12, kernel_size=(1, 64), stride=(1, 1))\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "nb_chev_filterstdn = 64\n",
    "nb_time_filterstdn = 64\n",
    "nb_blockstdn = 1\n",
    "dropouts=0.01\n",
    "Student=ASTGCN( nb_blockstdn, in_channels, K, nb_chev_filterstdn, nb_time_filterstdn, time_strides, num_for_predict, len_input, num_of_vertices,dropouts).to(DEVICE)\n",
    "print(Student)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0da0e88-db3d-4690-baaa-2d5f1939e756",
   "metadata": {},
   "source": [
    "# intialized teacher model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d119e88c-aab0-4bbf-84bd-f07093dd3974",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Teacher's state_dict:\n",
      "_blocklist.0._temporal_attention._U1 \t torch.Size([307]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._U2 \t torch.Size([1, 307]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._U3 \t torch.Size([1]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._be \t torch.Size([1, 12, 12]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._Ve \t torch.Size([12, 12]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._W1 \t torch.Size([12]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._W2 \t torch.Size([1, 12]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._W3 \t torch.Size([1]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._bs \t torch.Size([1, 307, 307]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._Vs \t torch.Size([307, 307]) \t cuda:0\n",
      "_blocklist.0._chebconv_attention._weight \t torch.Size([3, 1, 64]) \t cuda:0\n",
      "_blocklist.0._chebconv_attention._bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._time_convolution.weight \t torch.Size([64, 64, 1, 3]) \t cuda:0\n",
      "_blocklist.0._time_convolution.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._residual_convolution.weight \t torch.Size([64, 1, 1, 1]) \t cuda:0\n",
      "_blocklist.0._residual_convolution.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._layer_norm.weight \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._layer_norm.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._temporal_attention._U1 \t torch.Size([307]) \t cuda:0\n",
      "_blocklist.1._temporal_attention._U2 \t torch.Size([64, 307]) \t cuda:0\n",
      "_blocklist.1._temporal_attention._U3 \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._temporal_attention._be \t torch.Size([1, 12, 12]) \t cuda:0\n",
      "_blocklist.1._temporal_attention._Ve \t torch.Size([12, 12]) \t cuda:0\n",
      "_blocklist.1._spatial_attention._W1 \t torch.Size([12]) \t cuda:0\n",
      "_blocklist.1._spatial_attention._W2 \t torch.Size([64, 12]) \t cuda:0\n",
      "_blocklist.1._spatial_attention._W3 \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._spatial_attention._bs \t torch.Size([1, 307, 307]) \t cuda:0\n",
      "_blocklist.1._spatial_attention._Vs \t torch.Size([307, 307]) \t cuda:0\n",
      "_blocklist.1._chebconv_attention._weight \t torch.Size([3, 64, 64]) \t cuda:0\n",
      "_blocklist.1._chebconv_attention._bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._time_convolution.weight \t torch.Size([64, 64, 1, 3]) \t cuda:0\n",
      "_blocklist.1._time_convolution.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._residual_convolution.weight \t torch.Size([64, 64, 1, 1]) \t cuda:0\n",
      "_blocklist.1._residual_convolution.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._layer_norm.weight \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.1._layer_norm.bias \t torch.Size([64]) \t cuda:0\n",
      "_final_conv.weight \t torch.Size([12, 12, 1, 64]) \t cuda:0\n",
      "_final_conv.bias \t torch.Size([12]) \t cuda:0\n",
      "Teacher's total params: 450159\n",
      "Optimizer's state_dict:\n",
      "state \t {}\n",
      "param_groups \t [{'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 1e-06, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37]}]\n"
     ]
    }
   ],
   "source": [
    "#------------------------------------------------------\n",
    "learning_rate = 0.001\n",
    "# \n",
    "optimizerTeacher = optim.Adam(Teacher.parameters(), lr=1e-3,weight_decay=1e-6)#l2 regularization applied in weightdecay<0\n",
    "\n",
    "print('Teacher\\'s state_dict:')\n",
    "total_param = 0\n",
    "for param_tensor in Teacher.state_dict():\n",
    "    print(param_tensor, '\\t', Teacher.state_dict()[param_tensor].size(), '\\t', Teacher.state_dict()[param_tensor].device)\n",
    "    total_param += np.prod(Teacher.state_dict()[param_tensor].size())\n",
    "print('Teacher\\'s total params:', total_param)\n",
    "#--------------------------------------------------\n",
    "print('Optimizer\\'s state_dict:')\n",
    "for var_name in optimizerTeacher.state_dict():\n",
    "    print(var_name, '\\t', optimizerTeacher.state_dict()[var_name])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb1180f-7517-4d87-9982-dbd7c64860a9",
   "metadata": {},
   "source": [
    "## Intialized student model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c2ae6e1d-3714-4695-bbae-16112488590e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Teacher's state_dict:\n",
      "_blocklist.0._temporal_attention._U1 \t torch.Size([307]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._U2 \t torch.Size([1, 307]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._U3 \t torch.Size([1]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._be \t torch.Size([1, 12, 12]) \t cuda:0\n",
      "_blocklist.0._temporal_attention._Ve \t torch.Size([12, 12]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._W1 \t torch.Size([12]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._W2 \t torch.Size([1, 12]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._W3 \t torch.Size([1]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._bs \t torch.Size([1, 307, 307]) \t cuda:0\n",
      "_blocklist.0._spatial_attention._Vs \t torch.Size([307, 307]) \t cuda:0\n",
      "_blocklist.0._chebconv_attention._weight \t torch.Size([3, 1, 64]) \t cuda:0\n",
      "_blocklist.0._chebconv_attention._bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._time_convolution.weight \t torch.Size([64, 64, 1, 3]) \t cuda:0\n",
      "_blocklist.0._time_convolution.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._residual_convolution.weight \t torch.Size([64, 1, 1, 1]) \t cuda:0\n",
      "_blocklist.0._residual_convolution.bias \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._layer_norm.weight \t torch.Size([64]) \t cuda:0\n",
      "_blocklist.0._layer_norm.bias \t torch.Size([64]) \t cuda:0\n",
      "_final_conv.weight \t torch.Size([12, 12, 1, 64]) \t cuda:0\n",
      "_final_conv.bias \t torch.Size([12]) \t cuda:0\n",
      "Student's total params: 211518\n",
      "Optimizer's state_dict:\n",
      "state \t {}\n",
      "param_groups \t [{'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 1e-06, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]}]\n"
     ]
    }
   ],
   "source": [
    "#------------------------------------------------------\n",
    "# weight_decay=1e-3\n",
    "optimizerStudent = optim.Adam(Student.parameters(), lr=1e-3,weight_decay=1e-6)\n",
    "\n",
    "print('Teacher\\'s state_dict:')\n",
    "total_param = 0\n",
    "for param_tensor in Student.state_dict():\n",
    "    print(param_tensor, '\\t', Student.state_dict()[param_tensor].size(), '\\t', Student.state_dict()[param_tensor].device)\n",
    "    total_param += np.prod(Student.state_dict()[param_tensor].size())\n",
    "print('Student\\'s total params:', total_param)\n",
    "#--------------------------------------------------\n",
    "print('Optimizer\\'s state_dict:')\n",
    "for var_name in optimizerStudent.state_dict():\n",
    "    print(var_name, '\\t', optimizerStudent.state_dict()[var_name])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "99cdfc51-7445-4886-9dd8-78649487c0d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def masked_mae(preds, labels, null_val=np.nan):\n",
    "    if np.isnan(null_val):\n",
    "        mask = ~torch.isnan(labels)\n",
    "    else:\n",
    "        mask = (labels != null_val)\n",
    "    mask = mask.float()\n",
    "    mask /= torch.mean((mask))\n",
    "    mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)\n",
    "    loss = torch.abs(preds - labels)\n",
    "    loss = loss * mask\n",
    "    loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)\n",
    "    return torch.mean(loss)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "778ff538-0c64-4191-9752-f5fe61f5d047",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSELoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.mse = nn.MSELoss()\n",
    "        \n",
    "    def forward(self,ytil,y):\n",
    "        return torch.sqrt(self.mse(ytil,y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "142d4546-b1a4-41fa-9e86-7056d8016980",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha=0.4\n",
    "masked_flag=0\n",
    "criterionStudent=RMSELoss().to(DEVICE)\n",
    "\n",
    "criterion_masked = masked_mae\n",
    "loss_function = 'mae'\n",
    "\n",
    "metric_method = 'unmask'\n",
    "missing_value=0.0\n",
    "\n",
    "if loss_function=='masked_mse':\n",
    "    criterion_masked = masked_mse         #nn.MSELoss().to(DEVICE)\n",
    "    masked_flag=1\n",
    "elif loss_function=='masked_mae':\n",
    "    criterion_masked = masked_mae\n",
    "    masked_flag = 1\n",
    "elif loss_function == 'mae':\n",
    "    criterion = nn.L1Loss().to(DEVICE)\n",
    "    masked_flag = 0\n",
    "elif loss_function == 'rmse':\n",
    "    criterion = nn.MSELoss().to(DEVICE)\n",
    "    masked_flag= 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "77cace8d-c2a4-47f4-a9ef-0d6c197d9a46",
   "metadata": {},
   "outputs": [],
   "source": [
    "ValLoss=[]\n",
    "TrainLoss=[]\n",
    "distloss=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "06cc7bee-ed67-447c-b642-27303b61a7fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_val_loss_mstgcn(net, val_loader, criterion,  masked_flag,missing_value,sw, epoch, edge_index_data, limit=None):\n",
    "    '''\n",
    "    for rnn, compute mean loss on validation set\n",
    "    :param net: model\n",
    "    :param val_loader: torch.utils.data.utils.DataLoader\n",
    "    :param criterion: torch.nn.MSELoss\n",
    "    :param sw: tensorboardX.SummaryWriter\n",
    "    :param global_step: int, current global_step\n",
    "    :param limit: int,\n",
    "    :return: val_loss\n",
    "    '''\n",
    "    net.train(False)  # ensure dropout layers are in evaluation mode\n",
    "    with torch.no_grad():\n",
    "        val_loader_length = len(val_loader)  # nb of batch\n",
    "        tmp = []  # batch loss\n",
    "        for batch_index, batch_data in enumerate(val_loader):\n",
    "            encoder_inputs, labels = batch_data\n",
    "            outputs = net(encoder_inputs, edge_index_data)\n",
    "            if masked_flag:\n",
    "                loss = criterion(outputs, labels)\n",
    "            else:\n",
    "                loss = criterion(outputs, labels)\n",
    "            tmp.append(loss.item())\n",
    "            if batch_index % 20 == 0:\n",
    "                print('validation batch %s / %s, loss: %.2f' % (batch_index + 1, val_loader_length, loss.item()))\n",
    "            if (limit is not None) and batch_index >= limit:\n",
    "                break\n",
    "\n",
    "        validation_loss = sum(tmp) / len(tmp)\n",
    "        ValLoss.append(validation_loss)\n",
    "        sw.add_scalar('validation_loss', validation_loss, epoch)\n",
    "        print(\"rata rata val_loss\",validation_loss,\"pada epoch:\",epoch)\n",
    "    return validation_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "24a1642d-f5b7-4524-9ea5-a3856b045e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "global_step = 0\n",
    "best_epoch = 0\n",
    "best_val_loss = np.inf\n",
    "start_time= time()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f18b437e-8379-420d-a6a8-dea3f043807b",
   "metadata": {},
   "source": [
    "# load weights from teacher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "deb20d13-5500-4c21-8a83-371607e2b86e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ASTGCN(\n",
       "  (_blocklist): ModuleList(\n",
       "    (0): ASTGCNBlock(\n",
       "      (_temporal_attention): TemporalAttention()\n",
       "      (_spatial_attention): SpatialAttention()\n",
       "      (_chebconv_attention): ChebConvAttention(1, 64, K=3, normalization=None)\n",
       "      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))\n",
       "      (_residual_convolution): Conv2d(1, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
       "      (_dropout): Dropout2d(p=0.01, inplace=True)\n",
       "    )\n",
       "    (1): ASTGCNBlock(\n",
       "      (_temporal_attention): TemporalAttention()\n",
       "      (_spatial_attention): SpatialAttention()\n",
       "      (_chebconv_attention): ChebConvAttention(64, 64, K=3, normalization=None)\n",
       "      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))\n",
       "      (_residual_convolution): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
       "      (_dropout): Dropout2d(p=0.01, inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (_final_conv): Conv2d(12, 12, kernel_size=(1, 64), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epoch = 22\n",
    "params_filename = os.path.join('./Teacher/', 'checkpoint_%s.pth' % epoch)\n",
    "loaded_checkpoint=torch.load(params_filename)\n",
    "epoch=loaded_checkpoint['epoch']\n",
    "Teacher.load_state_dict(loaded_checkpoint['model_state'])\n",
    "optimizerTeacher.load_state_dict(loaded_checkpoint['optimizer_state'])\n",
    "Teacher.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9efb359b-392a-48fd-8262-0cba113e44dc",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OrderedDict([('_blocklist.0._temporal_attention._U1', tensor([ 5.5869e-01, -1.2616e-01,  1.7045e-01,  2.3812e-02,  4.8195e-01,\n",
      "         6.3830e-01,  9.7228e-02,  6.9284e-01,  6.2028e-02,  1.3427e-01,\n",
      "         5.9608e-02,  2.1886e-01, -2.9185e-02, -2.1210e-01,  4.1344e-01,\n",
      "         4.8402e-01,  4.7639e-01,  6.2355e-01, -3.7197e-01,  3.5150e-01,\n",
      "         2.9970e-01, -5.1058e-02, -4.0474e-02,  3.1314e-01,  3.2762e-02,\n",
      "         8.3597e-01,  3.0440e-01,  1.0644e-02,  7.0973e-02, -2.7030e-01,\n",
      "        -4.8443e-02,  3.3125e-01, -1.4869e-01,  4.1654e-01,  1.9244e-01,\n",
      "         4.4903e-01,  5.4413e-01,  1.9074e-01,  1.0784e+00, -1.2609e-01,\n",
      "         1.2166e-01,  1.0992e-01, -5.5038e-01, -3.9896e-01,  4.6273e-02,\n",
      "         7.7185e-01,  6.4420e-01,  5.1098e-01,  5.1247e-01, -1.4301e-01,\n",
      "        -5.7959e-01, -3.9534e-01,  1.8895e-01,  6.8797e-01,  3.2177e-01,\n",
      "        -2.0934e-01,  4.2424e-01,  4.4988e-03, -1.9095e-01, -1.6160e-02,\n",
      "         4.4068e-01, -1.2400e-01,  7.1368e-02,  7.2863e-02,  5.0593e-02,\n",
      "         7.4039e-03, -1.7763e-01, -1.7515e-01, -2.0903e-03,  3.7380e-01,\n",
      "         6.7138e-01, -3.8779e-01, -3.4406e-01,  1.6796e-01,  4.2206e-01,\n",
      "         1.4288e-01, -3.2706e-01,  2.5918e-01, -4.9453e-01, -3.1277e-02,\n",
      "         2.8246e-01,  5.7648e-02,  3.6464e-01,  2.4086e-01,  3.5653e-01,\n",
      "        -2.7593e-01,  4.8316e-01,  1.9958e-01, -5.2228e-03,  7.9464e-01,\n",
      "         4.7907e-02,  7.0998e-01,  3.7491e-01,  4.1578e-01, -1.4951e-01,\n",
      "        -1.5120e-01, -1.5648e-01, -5.6269e-01,  4.7437e-01,  3.6708e-02,\n",
      "        -6.3113e-01,  6.7417e-01,  2.6620e-01,  4.7005e-01, -7.2324e-02,\n",
      "         2.7280e-01, -1.7195e-01,  6.7529e-01,  5.5547e-01,  3.7387e-01,\n",
      "         2.1608e-01,  2.6955e-01,  3.6698e-01, -3.2378e-02, -1.2339e-02,\n",
      "         1.2952e+00,  5.3584e-01, -1.9901e-01,  2.1796e-01, -1.5588e-02,\n",
      "        -2.1095e-01, -1.2527e-01, -2.3581e-01, -4.2324e-02,  2.2225e-02,\n",
      "         2.2350e-01,  9.6348e-02,  2.6033e-02,  5.7140e-01,  2.3795e-01,\n",
      "         4.0746e-01, -3.4167e-02,  2.5358e-01,  5.2783e-01,  3.4058e-01,\n",
      "        -5.0794e-02, -2.8207e-01,  3.1183e-01, -1.1548e-01, -2.7182e-01,\n",
      "         4.3544e-01,  8.9026e-01,  9.6825e-02, -3.7350e-01,  3.1707e-01,\n",
      "         3.8431e-01,  7.7058e-01,  1.0474e-02,  3.9634e-01,  4.4636e-01,\n",
      "         2.4807e-01,  4.3869e-01,  1.3041e+00,  4.3592e-01,  5.5105e-01,\n",
      "         2.7191e-01,  7.6333e-02,  4.0662e-01,  9.7165e-01, -9.2994e-02,\n",
      "        -3.1480e-01,  4.8367e-01,  3.7643e-01,  7.0563e-01, -6.1520e-01,\n",
      "         4.1556e-01,  1.0982e-02,  4.1573e-01, -4.4382e-01, -7.4125e-02,\n",
      "         6.2130e-01,  7.5733e-01,  5.8916e-02,  4.2248e-01,  5.9407e-02,\n",
      "         4.6702e-01,  1.0756e-01,  1.9735e-01, -1.2205e-01, -7.2016e-02,\n",
      "         6.6643e-01,  3.4311e-01,  3.0658e-01,  5.6394e-01, -2.1105e-02,\n",
      "         1.4407e-01,  2.7421e-01,  8.9039e-01,  8.1869e-01,  3.8267e-01,\n",
      "         1.8767e-01,  5.7386e-02,  2.1118e-01,  5.7501e-02,  6.5651e-01,\n",
      "         6.5471e-01,  3.7125e-02,  4.8229e-01,  5.3760e-01,  3.4208e-01,\n",
      "         3.7476e-01,  5.6014e-01,  2.4840e-01,  2.8846e-01, -1.6918e-02,\n",
      "        -9.8674e-04,  2.8700e-01, -4.0133e-01,  1.3132e-01,  4.9350e-01,\n",
      "         2.1896e-01,  4.2004e-01,  5.0958e-02,  3.8918e-01,  4.5736e-01,\n",
      "         1.0850e-01,  4.9166e-01,  6.0341e-01,  4.9970e-03,  1.1674e-01,\n",
      "        -2.7736e-02, -1.2614e-01,  3.0550e-01,  3.6738e-01,  3.7199e-01,\n",
      "         4.1330e-02,  5.6088e-01,  2.4849e-01, -7.4590e-03,  2.2689e-01,\n",
      "         8.4032e-01,  6.4423e-01,  6.7745e-01,  3.1161e-01,  6.2641e-02,\n",
      "         3.9006e-01,  4.9790e-02,  4.8887e-01, -1.6085e-01,  4.0368e-01,\n",
      "        -6.1416e-01,  2.7563e-01,  2.6587e-01,  6.0389e-01,  4.5542e-03,\n",
      "         4.7685e-01,  2.4424e-01,  2.8541e-01,  2.0816e-01,  7.0309e-02,\n",
      "         5.5611e-01,  1.0623e-01, -9.2533e-02, -3.3332e-02,  2.1543e-01,\n",
      "        -2.4076e-01,  4.7464e-01,  1.1496e-01,  6.7287e-01, -2.1869e-01,\n",
      "         4.4235e-02,  2.6842e-01, -2.8355e-01, -6.5997e-02,  8.4556e-02,\n",
      "         1.6470e-01,  6.3343e-01, -4.4038e-02, -1.0828e-01, -3.0622e-01,\n",
      "        -5.8202e-02,  4.1097e-02,  8.9812e-02,  4.2567e-01,  6.3845e-02,\n",
      "         1.9012e-01,  5.6291e-01,  4.4944e-01,  7.1398e-01,  1.0170e+00,\n",
      "         2.5181e-01,  1.9002e-01,  5.5677e-01,  5.1676e-01,  3.3493e-01,\n",
      "         9.5235e-01,  2.9512e-01,  3.6035e-01, -2.0386e-01,  3.5213e-01,\n",
      "         1.3287e-01,  1.7972e-01,  3.7837e-01,  5.0887e-01, -4.6342e-03,\n",
      "         7.3437e-01,  1.7196e-01, -5.9274e-01,  1.9919e-01,  2.4792e-01,\n",
      "         6.6915e-01, -1.6167e-01, -9.6619e-01, -1.0853e-02,  8.0929e-02,\n",
      "        -1.4184e-01,  4.0890e-02], device='cuda:0')), ('_blocklist.0._temporal_attention._U2', tensor([[-0.0306,  0.3102, -0.4037,  0.2948,  0.0195, -0.1467,  0.3075, -0.0990,\n",
      "         -0.0150,  0.2254,  0.0914, -0.0997, -0.0353,  0.1212,  0.2071, -0.1646,\n",
      "         -0.0381,  0.0378,  0.0254,  0.0590,  0.0723,  0.1830,  0.4348,  0.1360,\n",
      "          0.3756, -0.0074,  0.0729,  0.1695,  0.1709,  0.0264,  0.0929, -0.0564,\n",
      "          0.0771, -0.0412, -0.0955, -0.0936,  0.2036, -0.2029, -0.0678,  0.1648,\n",
      "         -0.0120, -0.1348,  0.1722, -0.0382,  0.0758,  0.0498, -0.1347, -0.1389,\n",
      "         -0.3192, -0.1053, -0.1303,  0.2210, -0.0822,  0.0011,  0.1563,  0.0864,\n",
      "         -0.0590,  0.0621,  0.1341,  0.1778,  0.0545,  0.1037,  0.0374, -0.1218,\n",
      "         -0.1460, -0.0129,  0.0782,  0.0459,  0.1153,  0.0626,  0.0630,  0.1531,\n",
      "          0.0176, -0.0235,  0.0847, -0.1675,  0.0235, -0.0488,  0.1784, -0.1627,\n",
      "          0.1638, -0.1865,  0.1181, -0.1038, -0.2393,  0.1897,  0.1252,  0.0777,\n",
      "         -0.2935, -0.0251, -0.1492, -0.1377, -0.1869, -0.0271, -0.0921, -0.0401,\n",
      "          0.1328,  0.0744,  0.1506,  0.2327,  0.2133, -0.2587,  0.1298,  0.1794,\n",
      "          0.2247,  0.2269,  0.0181,  0.0240, -0.1568, -0.2051, -0.0314,  0.0915,\n",
      "          0.3286,  0.0455, -0.2776, -0.1982, -0.0657,  0.0035, -0.1472, -0.0841,\n",
      "         -0.0347, -0.0394,  0.0333, -0.1888,  0.0472,  0.1482,  0.1781,  0.1276,\n",
      "          0.0449,  0.0907,  0.0327,  0.1406,  0.1347,  0.0840,  0.3422,  0.2750,\n",
      "          0.1266, -0.0784, -0.1782,  0.2551,  0.0196, -0.2860, -0.0726,  0.2790,\n",
      "          0.1310,  0.0636, -0.2520, -0.1358,  0.2179,  0.0567, -0.0159, -0.1229,\n",
      "         -0.1254,  0.0362, -0.4921,  0.1460,  0.0138,  0.0406,  0.1298,  0.0872,\n",
      "          0.3216, -0.2908,  0.1330,  0.0862,  0.2083, -0.0165, -0.1717, -0.1246,\n",
      "          0.1499, -0.0085, -0.0102, -0.0662,  0.1686, -0.1663, -0.0270,  0.1215,\n",
      "         -0.0320,  0.0053, -0.1188,  0.0325,  0.1106,  0.1680,  0.0069,  0.1044,\n",
      "          0.1872,  0.0743,  0.0022, -0.1978,  0.0046,  0.1754,  0.1300,  0.1194,\n",
      "          0.1237,  0.1748,  0.1609,  0.1766,  0.0628,  0.1221,  0.1236,  0.1337,\n",
      "          0.1308, -0.1768, -0.0923,  0.1704, -0.2352,  0.0580, -0.0288,  0.0739,\n",
      "          0.0245, -0.1262, -0.2166, -0.0796, -0.0583,  0.0218, -0.0973, -0.0839,\n",
      "          0.1217, -0.2400,  0.0646, -0.2931, -0.0607, -0.0471, -0.0110,  0.1066,\n",
      "          0.0790,  0.0849, -0.0631, -0.1511,  0.1193,  0.0666,  0.1087, -0.1641,\n",
      "         -0.3485,  0.0762,  0.0055,  0.0761,  0.0888, -0.1267,  0.0637, -0.1531,\n",
      "          0.1162, -0.3754, -0.0400, -0.0561,  0.1396, -0.5122, -0.1004,  0.3353,\n",
      "         -0.0177,  0.1309,  0.0339, -0.1755,  0.0719, -0.0664, -0.1565,  0.1313,\n",
      "          0.1249, -0.0449, -0.1435,  0.2161,  0.1003, -0.0392,  0.0405, -0.0328,\n",
      "          0.0553, -0.0424,  0.2107,  0.0174, -0.3548,  0.0208,  0.1757, -0.1976,\n",
      "         -0.0818,  0.1073,  0.0860,  0.0567, -0.0608, -0.2483, -0.1152, -0.2205,\n",
      "         -0.1558, -0.0382, -0.1882,  0.0556, -0.2088,  0.1728,  0.0845, -0.0706,\n",
      "          0.1934,  0.0647, -0.1990, -0.1213, -0.1547,  0.0288, -0.1060,  0.0349,\n",
      "          0.2259,  0.0044, -0.2274, -0.0024, -0.0663,  0.1798,  0.1620,  0.1498,\n",
      "         -0.0013,  0.1434, -0.0795]], device='cuda:0')), ('_blocklist.0._temporal_attention._U3', tensor([0.0911], device='cuda:0')), ('_blocklist.0._temporal_attention._be', tensor([[[ 0.4051,  0.4769, -0.7119,  0.1629, -0.5095,  0.1893,  0.2694,\n",
      "           0.4019, -0.4006, -0.5413,  0.3358, -0.6520],\n",
      "         [ 0.1980,  0.1596, -0.4733,  0.1235, -0.1749,  0.2037,  0.1491,\n",
      "           0.1410, -0.2756, -0.2141,  0.1193, -0.3609],\n",
      "         [ 0.1971,  0.1674, -0.2817,  0.1816,  0.0820,  0.2388,  0.1361,\n",
      "           0.3135, -0.3276, -0.2896,  0.2002, -0.2070],\n",
      "         [ 0.6961,  0.6000, -0.7649,  0.3953, -0.3645,  0.3591,  0.3408,\n",
      "           0.5327, -0.4811, -0.5616,  0.4439, -0.5961],\n",
      "         [ 0.3973,  0.5018, -0.6493,  0.2343, -0.4994,  0.1955,  0.3254,\n",
      "           0.3389, -0.4098, -0.4576,  0.2789, -0.6014],\n",
      "         [ 0.5030,  0.5028, -0.5470,  0.2378, -0.4396,  0.4356,  0.3522,\n",
      "           0.3736, -0.4189, -0.4687,  0.4156, -0.5571],\n",
      "         [ 0.9242,  0.9750, -1.0616,  0.7358, -0.2898,  0.8155,  0.8288,\n",
      "           0.9433, -0.9287, -0.9413,  0.7999, -1.0852],\n",
      "         [ 1.0012,  0.9723, -0.9750,  0.7997, -0.0971,  0.8017,  0.7826,\n",
      "           0.9189, -0.8046, -0.8308,  0.9348, -1.0404],\n",
      "         [ 0.8105,  0.8203, -0.7949,  0.7608, -0.0471,  0.7038,  0.6919,\n",
      "           0.6057, -0.6050, -0.5833,  0.7560, -0.8306],\n",
      "         [ 0.5107,  0.4437, -0.5316,  0.4585,  0.1236,  0.4726,  0.3561,\n",
      "           0.2040, -0.3372, -0.2725,  0.4111, -0.4566],\n",
      "         [ 0.7757,  0.7751, -0.8916,  0.6571, -0.1143,  0.5664,  0.5578,\n",
      "           0.6236, -0.6527, -0.6700,  0.6526, -0.7711],\n",
      "         [ 0.6344,  0.6525, -0.7359,  0.5136, -0.3586,  0.3768,  0.3953,\n",
      "           0.5863, -0.4722, -0.5517,  0.4666, -0.6556]]], device='cuda:0')), ('_blocklist.0._temporal_attention._Ve', tensor([[-0.2491, -0.5403, -0.2798,  0.0151, -0.2392, -0.0357, -0.5712, -1.0435,\n",
      "         -0.5701, -0.5834, -0.8684, -0.3983],\n",
      "        [-0.0578, -0.0518, -0.1040, -0.4354, -0.0407, -0.2706, -0.6547, -0.6543,\n",
      "         -0.6712, -0.5636, -0.6196, -0.4716],\n",
      "        [ 0.1092, -0.2420, -0.0029,  0.0446, -0.1596, -0.3568, -0.2895, -0.2723,\n",
      "         -0.3918, -0.4564, -0.5172, -0.1790],\n",
      "        [-0.1389,  0.3742,  0.3230,  0.3280,  0.1048,  0.3283, -0.0405, -0.2340,\n",
      "         -0.5715, -0.0012, -0.1885, -0.5254],\n",
      "        [-0.0916,  0.2907,  0.1067,  0.3358,  0.2053, -0.1653,  0.0902, -0.1907,\n",
      "         -0.2519, -0.0507, -0.1117, -0.2600],\n",
      "        [-0.1432,  0.2441,  0.3976,  0.0477,  0.1175,  0.3030, -0.3741, -0.2503,\n",
      "          0.2379,  0.2736, -0.0641, -0.2065],\n",
      "        [-0.1561,  0.2452, -0.1600,  0.1674,  0.2113,  0.2081, -0.0511,  0.0470,\n",
      "         -0.2104,  0.1395,  0.0410,  0.1695],\n",
      "        [-0.1849,  0.0427, -0.2649,  0.3086,  0.3991,  0.4777,  0.2403, -0.1718,\n",
      "          0.1719,  0.4193,  0.0179, -0.2830],\n",
      "        [ 0.3634,  0.3029, -0.3082, -0.0762,  0.3452,  0.4470,  0.1264,  0.0491,\n",
      "          0.0933,  0.0587,  0.2257,  0.2161],\n",
      "        [ 0.0432,  0.2592,  0.3465,  0.1267, -0.2960,  0.2329,  0.1298,  0.3237,\n",
      "         -0.0037,  0.3595,  0.2484,  0.1119],\n",
      "        [ 0.3902, -0.1585,  0.1946,  0.1527,  0.2755, -0.2145,  0.2446,  0.0924,\n",
      "          0.4221, -0.0557,  0.0452,  0.4686],\n",
      "        [ 0.2590,  0.0325, -0.1975,  0.3712,  0.2873,  0.2946, -0.0348,  0.0288,\n",
      "          0.0272,  0.1797,  0.3111,  0.5178]], device='cuda:0')), ('_blocklist.0._spatial_attention._W1', tensor([0.2056, 0.0747, 0.0569, 0.7383, 0.3440, 0.7446, 0.4611, 0.0177, 0.0071,\n",
      "        0.1593, 0.6364, 0.1854], device='cuda:0')), ('_blocklist.0._spatial_attention._W2', tensor([[-0.4221, -0.3420,  0.4687, -0.2408,  0.0045, -0.5276, -0.4050, -0.5053,\n",
      "          0.1763,  0.3147, -0.3464,  0.3699]], device='cuda:0')), ('_blocklist.0._spatial_attention._W3', tensor([0.8064], device='cuda:0')), ('_blocklist.0._spatial_attention._bs', tensor([[[ 2.4766e-03, -6.3336e-02, -2.4966e-01,  ..., -1.6666e-03,\n",
      "           3.1818e-04,  3.2685e-04],\n",
      "         [ 1.0979e-03, -4.2330e-03,  7.0305e-01,  ..., -1.2568e-03,\n",
      "          -2.8828e-04,  4.2418e-04],\n",
      "         [ 3.2806e-02, -8.6569e-02, -6.8649e-01,  ...,  1.6346e-02,\n",
      "           1.9564e-03, -3.6140e-03],\n",
      "         ...,\n",
      "         [ 5.6611e-03, -1.8504e-02, -4.4438e-01,  ..., -1.3936e-02,\n",
      "          -4.0257e-04,  9.3663e-05],\n",
      "         [ 1.1568e-03, -1.9817e-02, -3.1491e-01,  ..., -8.0210e-04,\n",
      "          -5.0028e-05,  3.1136e-04],\n",
      "         [-2.4992e-03, -1.8315e-02, -5.7170e-01,  ..., -1.9651e-05,\n",
      "           3.3167e-05,  2.0951e-04]]], device='cuda:0')), ('_blocklist.0._spatial_attention._Vs', tensor([[ 0.0660, -0.0253, -0.0748,  ...,  0.0110,  0.0257, -0.0049],\n",
      "        [-0.0630, -0.1067, -0.0059,  ..., -0.0389, -0.0462, -0.0812],\n",
      "        [ 0.4279,  0.1720, -0.7636,  ..., -0.1273,  0.0775,  0.0991],\n",
      "        ...,\n",
      "        [-0.0357, -0.1385, -0.0500,  ..., -0.1110, -0.0469, -0.0582],\n",
      "        [-0.0287, -0.0542, -0.1038,  ..., -0.0216,  0.0046, -0.0172],\n",
      "        [-0.0264, -0.0381, -0.0922,  ..., -0.0195,  0.0068, -0.0135]],\n",
      "       device='cuda:0')), ('_blocklist.0._chebconv_attention._weight', tensor([[[-1.1374,  1.3441,  0.4570,  0.2338, -1.2162,  1.3141,  0.6754,\n",
      "           1.1854,  0.8001,  0.9296,  0.5508, -0.7800, -2.5263, -0.7966,\n",
      "          -0.9285, -1.4546, -2.0875,  0.8660,  1.1321, -0.6935, -1.0920,\n",
      "          -1.1286, -1.6515,  0.4174,  1.2573, -0.9668, -1.4409,  0.7205,\n",
      "           0.8696, -1.1616,  1.0054,  0.8439,  1.2548,  1.4334, -1.0448,\n",
      "          -1.3743,  0.3762,  1.3187, -1.0260, -1.0030,  1.1413,  0.4861,\n",
      "           0.7392, -0.7763, -1.6843,  1.0560, -0.4663,  0.8272, -1.3094,\n",
      "           0.4763,  0.1039, -0.7607, -0.5213, -1.3109, -0.6852, -1.6706,\n",
      "          -0.8979, -0.3935, -1.0196, -0.3994, -0.5840, -2.1946,  1.4726,\n",
      "          -0.6094]],\n",
      "\n",
      "        [[ 0.6722, -0.7243, -0.5613,  0.8534,  0.6731, -0.9829, -0.4383,\n",
      "          -0.6323,  0.0357, -0.3947, -0.4272,  1.2509,  0.7343,  0.3833,\n",
      "           0.5969,  0.3753,  0.9576, -0.6371, -0.7016,  0.5606,  0.5534,\n",
      "          -0.0977,  0.6884, -0.2935, -0.9990,  0.1729,  0.4327, -0.2600,\n",
      "          -0.3048,  0.6509, -0.6513, -0.2910, -1.1507, -1.5577,  0.7228,\n",
      "           0.6382, -0.0128, -0.7325,  0.6454,  0.4195, -0.2571, -0.3970,\n",
      "          -0.8047,  0.7708,  0.8765, -0.5974, -0.1629,  0.3962,  0.5860,\n",
      "           0.3466, -0.7255,  0.4234, -0.1944,  0.5046,  1.3516,  0.5505,\n",
      "           0.3576,  0.6413,  0.7393,  0.4263,  0.1906,  0.9354, -0.9212,\n",
      "           1.3674]],\n",
      "\n",
      "        [[ 1.0063, -1.2444, -0.4197, -0.0615,  1.0460, -1.1487, -0.7938,\n",
      "          -0.9126, -0.6519, -0.5960, -0.3340,  0.4392,  2.0117,  0.6968,\n",
      "           0.6673,  1.0484,  1.9449, -0.6895, -1.0272,  0.4359,  1.0307,\n",
      "           0.7863,  1.1149, -0.3548, -1.0637,  0.6173,  1.1042, -0.7950,\n",
      "          -0.9497,  1.2295, -0.9659, -0.9182, -1.1818, -1.1482,  1.0699,\n",
      "           1.1742, -0.3191, -1.2617,  1.1081,  0.8527, -0.9335, -0.5052,\n",
      "          -0.6905,  0.6342,  1.4484, -1.0215,  0.2561, -0.7766,  1.2032,\n",
      "          -0.4090, -0.3302,  0.5782,  0.6514,  0.9225,  0.6735,  1.1883,\n",
      "           0.6729,  0.0475,  0.9222,  0.4727,  0.4593,  1.6129, -1.2743,\n",
      "           0.4607]]], device='cuda:0')), ('_blocklist.0._chebconv_attention._bias', tensor([ 0.1896,  0.3385, -0.0627,  0.4380,  0.0537,  0.6546,  0.5898,  0.8050,\n",
      "         0.2824,  0.0102,  0.2665,  0.6960,  0.0238,  0.6797,  0.2680, -0.0044,\n",
      "         0.0488,  0.1618,  0.1223, -0.0035,  0.7999, -0.0016, -0.0032, -0.0121,\n",
      "         0.1471, -0.0043, -0.0032,  0.4602,  0.0471, -0.0238,  0.3500,  0.8591,\n",
      "         0.7863,  0.0238, -0.0225,  0.2066, -0.0021,  0.2852,  0.4691,  0.1261,\n",
      "         0.6693,  0.4300,  0.0790,  0.2820,  0.0209,  0.1227,  0.0012,  0.2698,\n",
      "         0.1986,  0.5098,  0.1929,  0.4628, -0.0029,  0.0114,  0.2714,  0.0099,\n",
      "         0.0209,  0.6113,  0.3199,  0.2526, -0.0028,  0.0230,  0.0358,  0.6121],\n",
      "       device='cuda:0')), ('_blocklist.0._time_convolution.weight', tensor([[[[ 1.2437e-02,  4.4541e-02, -3.9087e-02]],\n",
      "\n",
      "         [[ 5.3268e-02,  2.5817e-02,  5.3906e-02]],\n",
      "\n",
      "         [[-1.1976e-01, -3.1276e-01, -6.8669e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-2.0905e-02,  1.7042e-01, -9.4659e-02]],\n",
      "\n",
      "         [[ 1.5866e-01, -2.7673e-02,  2.2686e-01]],\n",
      "\n",
      "         [[ 1.7270e-01,  9.8788e-03, -1.9030e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 9.7544e-02, -5.4296e-02, -2.1217e-02]],\n",
      "\n",
      "         [[-4.3512e-03, -1.0822e-02, -8.4501e-02]],\n",
      "\n",
      "         [[ 2.9036e-02, -7.0061e-02,  1.5525e-01]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 3.3043e-02, -6.3864e-02,  1.0839e-01]],\n",
      "\n",
      "         [[-2.6989e-01, -4.1256e-02,  3.0604e-02]],\n",
      "\n",
      "         [[ 1.2005e-02,  1.9804e-02,  4.0600e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.0771e-01,  5.2263e-02, -6.1128e-02]],\n",
      "\n",
      "         [[-6.4539e-03,  8.5135e-02,  1.0076e-01]],\n",
      "\n",
      "         [[-1.7681e-01, -1.9344e-01, -2.5819e-01]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-1.3618e-01, -1.3658e-01, -1.6126e-01]],\n",
      "\n",
      "         [[ 1.7481e-01,  1.2646e-01, -3.4917e-04]],\n",
      "\n",
      "         [[ 9.1005e-02, -6.1866e-02,  1.0257e-01]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 9.2560e-04, -5.3985e-02, -1.1761e-02]],\n",
      "\n",
      "         [[-1.2076e-02, -1.4461e-02, -6.7440e-02]],\n",
      "\n",
      "         [[-7.4921e-07, -2.0892e-06, -4.5196e-06]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-4.5114e-02, -4.4419e-02, -2.2785e-02]],\n",
      "\n",
      "         [[-1.2372e-02, -6.1680e-03, -2.4923e-03]],\n",
      "\n",
      "         [[-1.0056e-01, -5.2225e-02, -1.8447e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 9.9390e-02, -1.3901e-03,  3.1442e-02]],\n",
      "\n",
      "         [[-6.4639e-02, -1.6247e-02,  9.3836e-02]],\n",
      "\n",
      "         [[-1.6583e-01,  1.5622e-02,  1.5061e-01]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-5.9707e-02, -9.5721e-02,  6.8597e-02]],\n",
      "\n",
      "         [[-3.9090e-02, -1.0413e-01,  1.6298e-01]],\n",
      "\n",
      "         [[ 2.1606e-02,  4.1107e-02, -1.0743e-01]]],\n",
      "\n",
      "\n",
      "        [[[-1.7884e-02, -2.1417e-02, -1.8272e-04]],\n",
      "\n",
      "         [[-1.4062e-39, -9.6125e-04, -1.0560e-03]],\n",
      "\n",
      "         [[-5.6653e-39,  1.6204e-39, -5.1357e-39]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-2.1016e-02, -2.1012e-02, -3.3185e-16]],\n",
      "\n",
      "         [[ 7.9519e-40, -2.1476e-26, -2.3905e-26]],\n",
      "\n",
      "         [[-1.8909e-02, -2.9302e-02, -5.7768e-03]]]], device='cuda:0')), ('_blocklist.0._time_convolution.bias', tensor([ 0.5495,  0.7661,  0.2914,  0.3701,  0.8442,  0.1311,  0.4119,  0.3828,\n",
      "         0.4926,  0.6174,  0.4535,  0.7327,  0.0845,  0.6690,  0.0784,  0.6294,\n",
      "         0.5559,  0.5051,  0.7976,  0.0169,  0.7171,  0.7152,  0.7490,  0.5521,\n",
      "         0.6830,  0.8775,  0.8089,  0.8081,  0.1009,  0.6230,  0.5845,  0.1681,\n",
      "         0.9284,  0.8510,  0.6646,  0.3778,  0.4833,  0.5606,  0.8554,  0.8772,\n",
      "         0.3394, -0.0360,  0.8954, -0.0181, -0.0013,  0.0140,  0.4326,  0.3093,\n",
      "         0.7505,  0.7731, -0.0070,  0.2148,  0.8037,  0.2284,  0.6589,  0.1088,\n",
      "         0.2778,  0.9760,  0.1143,  0.1507,  0.9469,  0.5472,  0.6532,  0.0240],\n",
      "       device='cuda:0')), ('_blocklist.0._residual_convolution.weight', tensor([[[[-0.2945]]],\n",
      "\n",
      "\n",
      "        [[[ 0.6081]]],\n",
      "\n",
      "\n",
      "        [[[-0.4721]]],\n",
      "\n",
      "\n",
      "        [[[-0.3064]]],\n",
      "\n",
      "\n",
      "        [[[ 0.7092]]],\n",
      "\n",
      "\n",
      "        [[[-0.0327]]],\n",
      "\n",
      "\n",
      "        [[[-1.1402]]],\n",
      "\n",
      "\n",
      "        [[[-0.0595]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0557]]],\n",
      "\n",
      "\n",
      "        [[[-0.3556]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0530]]],\n",
      "\n",
      "\n",
      "        [[[-0.2666]]],\n",
      "\n",
      "\n",
      "        [[[-0.5946]]],\n",
      "\n",
      "\n",
      "        [[[ 0.3543]]],\n",
      "\n",
      "\n",
      "        [[[ 0.7631]]],\n",
      "\n",
      "\n",
      "        [[[ 0.3863]]],\n",
      "\n",
      "\n",
      "        [[[-0.5576]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0027]]],\n",
      "\n",
      "\n",
      "        [[[-0.8682]]],\n",
      "\n",
      "\n",
      "        [[[-0.0258]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1420]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0777]]],\n",
      "\n",
      "\n",
      "        [[[-0.2064]]],\n",
      "\n",
      "\n",
      "        [[[-0.6813]]],\n",
      "\n",
      "\n",
      "        [[[ 0.6787]]],\n",
      "\n",
      "\n",
      "        [[[-0.0388]]],\n",
      "\n",
      "\n",
      "        [[[ 0.8994]]],\n",
      "\n",
      "\n",
      "        [[[ 0.5230]]],\n",
      "\n",
      "\n",
      "        [[[-0.7244]]],\n",
      "\n",
      "\n",
      "        [[[-0.1605]]],\n",
      "\n",
      "\n",
      "        [[[-0.2483]]],\n",
      "\n",
      "\n",
      "        [[[-0.5776]]],\n",
      "\n",
      "\n",
      "        [[[-0.2177]]],\n",
      "\n",
      "\n",
      "        [[[ 0.6119]]],\n",
      "\n",
      "\n",
      "        [[[-0.7952]]],\n",
      "\n",
      "\n",
      "        [[[ 0.6040]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1135]]],\n",
      "\n",
      "\n",
      "        [[[ 0.9540]]],\n",
      "\n",
      "\n",
      "        [[[ 0.9468]]],\n",
      "\n",
      "\n",
      "        [[[ 0.6561]]],\n",
      "\n",
      "\n",
      "        [[[-1.0871]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0415]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1180]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0048]]],\n",
      "\n",
      "\n",
      "        [[[-0.0261]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0020]]],\n",
      "\n",
      "\n",
      "        [[[ 0.2001]]],\n",
      "\n",
      "\n",
      "        [[[-0.0095]]],\n",
      "\n",
      "\n",
      "        [[[ 0.5850]]],\n",
      "\n",
      "\n",
      "        [[[-0.4799]]],\n",
      "\n",
      "\n",
      "        [[[ 0.5382]]],\n",
      "\n",
      "\n",
      "        [[[-0.1649]]],\n",
      "\n",
      "\n",
      "        [[[ 0.5116]]],\n",
      "\n",
      "\n",
      "        [[[ 0.2674]]],\n",
      "\n",
      "\n",
      "        [[[-0.7592]]],\n",
      "\n",
      "\n",
      "        [[[-0.1154]]],\n",
      "\n",
      "\n",
      "        [[[-0.4384]]],\n",
      "\n",
      "\n",
      "        [[[ 0.8939]]],\n",
      "\n",
      "\n",
      "        [[[-0.2386]]],\n",
      "\n",
      "\n",
      "        [[[ 0.5699]]],\n",
      "\n",
      "\n",
      "        [[[-0.5864]]],\n",
      "\n",
      "\n",
      "        [[[-0.1796]]],\n",
      "\n",
      "\n",
      "        [[[ 0.8749]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0115]]]], device='cuda:0')), ('_blocklist.0._residual_convolution.bias', tensor([ 1.8396e-01,  2.1451e-01,  2.4532e-01,  1.3345e-01,  7.8418e-01,\n",
      "         1.7484e-01, -2.4895e-02,  7.8545e-01,  9.2629e-01,  1.7906e-01,\n",
      "         9.8112e-01,  5.2528e-01,  2.4358e-01,  6.8355e-01,  9.7541e-01,\n",
      "         2.4565e-01,  8.5844e-01,  7.0367e-01,  4.9858e-01,  6.7535e-01,\n",
      "         9.0274e-01,  7.9986e-01,  6.8831e-01,  4.2741e-01,  6.9267e-01,\n",
      "         7.7392e-01,  8.9835e-01,  9.3083e-01,  8.6462e-01,  7.5462e-04,\n",
      "         3.2514e-01,  3.9118e-01,  4.4428e-01,  2.8797e-01,  5.3871e-01,\n",
      "         6.3967e-01,  5.4795e-01,  1.1312e-01,  4.2926e-01,  9.2619e-01,\n",
      "         4.6711e-01,  1.8342e-03,  8.5991e-01,  1.3763e-01, -3.1927e-03,\n",
      "         1.1411e-02,  2.5153e-01,  3.7786e-01,  1.7062e-01,  8.6490e-01,\n",
      "         9.3059e-01,  5.8894e-01,  8.8679e-01,  1.6067e-01,  5.6725e-02,\n",
      "         9.3164e-01,  6.3354e-01,  8.1150e-01,  9.4908e-01,  5.7381e-01,\n",
      "         7.1210e-01,  2.1683e-01,  7.3618e-01,  5.1756e-02], device='cuda:0')), ('_blocklist.0._layer_norm.weight', tensor([ 0.2870,  0.1529,  0.7772,  0.6311,  0.7803,  0.0159,  0.8129,  0.7553,\n",
      "         0.0148,  0.7581,  0.0213,  0.2521,  0.4663,  0.2189,  0.8252,  0.7896,\n",
      "         0.8631, -0.0269,  0.4861,  0.7068,  0.2685, -0.0072,  0.8061,  1.0602,\n",
      "         0.8366,  0.0596,  0.9708,  0.5508,  0.3499,  0.1702,  0.0655,  0.4973,\n",
      "        -0.0057,  0.4688,  0.3073,  0.9971,  0.2158,  0.6690,  1.0097,  0.3036,\n",
      "         1.2256,  0.7091,  0.3236, -0.0159, -0.1167,  0.8193,  0.2158,  0.8824,\n",
      "         0.4056,  0.7745,  0.6473,  0.4543,  1.1908,  0.4839,  0.8259,  0.4673,\n",
      "         0.5476,  0.8466,  0.1154,  0.3506,  0.9437,  0.6112,  0.4852,  0.4879],\n",
      "       device='cuda:0')), ('_blocklist.0._layer_norm.bias', tensor([ 0.6871,  0.2270,  0.4828,  0.3722,  0.3155,  0.4797,  0.4597,  1.0043,\n",
      "         0.3702,  0.4096,  0.4186,  0.9145,  0.2988,  0.3837,  0.4256, -0.0212,\n",
      "         0.2176,  0.9407,  0.2269,  0.7173,  0.0802,  0.9014,  0.7320,  0.8541,\n",
      "         0.7128,  0.7018,  0.7346,  0.8430,  0.0315,  0.2985,  0.4149,  0.0869,\n",
      "         0.5008,  0.9368,  0.3413,  0.0399,  0.9533,  0.6373,  0.2278,  0.7280,\n",
      "         0.6343,  0.5344,  0.6105,  0.4882,  0.3201,  0.0482,  0.6136,  0.6783,\n",
      "         0.0506,  0.5924,  0.4220,  0.0977,  0.4263,  0.0239,  0.9330,  0.9796,\n",
      "         0.8476,  0.6143,  0.6920,  0.6581,  0.6842,  0.2507,  0.6669,  0.4959],\n",
      "       device='cuda:0')), ('_blocklist.1._temporal_attention._U1', tensor([-5.8005e-02, -5.6541e-02, -5.7476e-02,  2.4781e-03, -1.1654e-02,\n",
      "        -5.8521e-02, -1.6558e-03, -1.0804e-02, -5.6669e-02,  4.2195e-03,\n",
      "         2.2039e-04,  1.0436e-02, -5.9690e-02, -1.3117e-03, -7.3251e-04,\n",
      "        -4.6287e-02, -6.0353e-02, -5.1124e-02,  3.0118e-03, -6.6553e-03,\n",
      "        -4.6129e-02, -1.5297e-02, -7.3659e-03, -2.6278e-03,  8.5630e-05,\n",
      "        -5.6943e-02,  4.9192e-04,  2.3016e-03, -9.9333e-03, -5.7440e-02,\n",
      "        -1.0003e-02,  8.8875e-03,  3.3454e-03, -2.4321e-03, -5.6321e-02,\n",
      "        -4.4654e-02,  1.8234e-03, -5.8855e-02, -5.6672e-02, -5.7641e-02,\n",
      "        -2.0762e-02, -1.7745e-02, -5.9186e-02, -1.5765e-02, -4.4791e-03,\n",
      "        -5.1634e-02,  9.0514e-04, -2.7833e-02, -5.4486e-02, -4.0561e-03,\n",
      "         7.7685e-04, -6.7033e-03, -5.5539e-02, -5.7603e-02,  1.0367e-03,\n",
      "        -5.4281e-03, -6.0148e-02, -6.5005e-02, -6.3870e-02,  1.4773e-03,\n",
      "        -5.2516e-03, -8.8817e-03, -1.9415e-02, -5.1032e-02, -1.1720e-02,\n",
      "        -4.6724e-02,  9.4395e-04, -5.9035e-02, -7.6802e-03, -1.8038e-02,\n",
      "         2.9524e-03,  7.4652e-03, -5.9495e-02, -5.4057e-02, -3.4787e-02,\n",
      "        -3.9277e-03, -5.0918e-04, -5.5766e-02,  4.7421e-02, -5.7990e-02,\n",
      "        -1.1003e-02, -5.6637e-02, -1.8707e-02, -4.3261e-02, -5.9016e-02,\n",
      "        -8.1376e-03, -6.2545e-03, -3.3969e-02,  5.7514e-02, -4.3736e-03,\n",
      "        -5.6984e-02, -5.5975e-02, -6.2317e-02, -5.5294e-02, -4.3451e-03,\n",
      "        -6.0750e-02, -6.1929e-02, -1.7000e-02,  1.5613e-03, -5.8126e-02,\n",
      "         1.8494e-03, -1.0112e-02, -5.7210e-02, -5.8516e-02,  6.1358e-04,\n",
      "         1.4744e-03, -5.8255e-02, -1.7437e-03,  1.2891e-02, -4.1786e-04,\n",
      "         4.2371e-02, -1.5112e-02, -2.3172e-03, -1.3561e-02, -5.7612e-02,\n",
      "         1.8607e-03,  2.0266e-03, -6.2196e-02, -6.7613e-02, -6.5088e-02,\n",
      "        -1.4970e-02, -2.2822e-03, -5.9982e-02, -4.6677e-02, -5.9215e-02,\n",
      "        -5.7091e-02, -4.3897e-02, -6.0817e-02, -2.9406e-04, -1.0992e-02,\n",
      "        -1.9473e-02, -1.9632e-03, -1.2239e-03, -3.2506e-03, -2.2225e-02,\n",
      "         1.3387e-03,  2.0677e-03,  1.3162e-02, -5.5937e-02, -1.1231e-03,\n",
      "        -1.1626e-03, -1.8691e-04, -1.7418e-02, -1.1285e-02, -9.2093e-06,\n",
      "        -4.6753e-03, -5.9298e-02, -5.5726e-02, -3.2728e-03,  1.5421e-03,\n",
      "         1.4163e-03, -5.7893e-02, -1.1904e-02, -2.1228e-04, -6.0949e-02,\n",
      "        -1.4903e-02, -7.8257e-03, -4.4612e-02, -5.6923e-02, -5.8517e-02,\n",
      "        -1.0249e-02, -5.7144e-02, -5.9824e-03, -5.7923e-02, -4.8167e-05,\n",
      "        -3.8093e-02, -3.9661e-02, -5.6235e-02, -3.8042e-02,  1.3928e-03,\n",
      "         1.9934e-03, -5.6102e-02, -7.0629e-06, -5.4614e-02, -5.8662e-02,\n",
      "         2.0695e-03, -5.6790e-04,  3.1158e-03,  7.5422e-04,  7.4172e-04,\n",
      "        -9.1107e-06,  2.6334e-05, -2.8045e-02, -9.2093e-06, -1.0280e-02,\n",
      "        -2.1237e-03, -1.4531e-02, -2.6398e-02, -7.3332e-03, -7.4966e-06,\n",
      "        -5.9557e-03,  4.0440e-06,  2.7587e-05,  2.4421e-05,  1.3180e-05,\n",
      "         1.5695e-05, -1.0245e-03, -8.0386e-06, -9.2112e-06, -8.8017e-06,\n",
      "        -9.1926e-06, -5.2336e-02, -5.7027e-02, -6.0483e-02, -5.8657e-02,\n",
      "        -2.4332e-04, -1.0703e-02, -6.7194e-02, -6.0188e-02, -1.3994e-02,\n",
      "        -5.7940e-02, -5.7820e-02, -1.1979e-03,  4.2768e-04,  9.7113e-03,\n",
      "         2.4500e-03,  3.3110e-03,  9.6574e-04,  5.3108e-03, -8.7383e-04,\n",
      "        -4.6908e-02, -5.7077e-02, -5.6859e-02, -4.5867e-03, -1.5671e-02,\n",
      "         1.0478e-04, -6.5905e-03, -5.3617e-02, -6.1888e-04, -1.5097e-02,\n",
      "        -1.6809e-02,  8.2058e-04,  2.3254e-03, -5.2184e-02, -1.9076e-02,\n",
      "        -5.7929e-02, -5.9280e-02, -3.3753e-02, -1.8647e-04, -6.3003e-02,\n",
      "         3.0023e-03, -5.9008e-02, -4.2611e-02, -1.1948e-02,  6.3990e-04,\n",
      "        -5.3133e-02, -5.6951e-03,  2.0953e-03, -2.2774e-02, -2.4831e-03,\n",
      "        -5.3853e-02, -5.7761e-02, -5.8817e-02, -1.8116e-02,  2.6916e-03,\n",
      "        -2.3597e-03, -5.4761e-02, -9.0603e-04, -2.1788e-02, -5.3553e-02,\n",
      "        -5.8189e-02,  1.4026e-03, -1.2238e-02, -5.7746e-02, -2.0874e-03,\n",
      "        -4.8820e-02, -1.9290e-03, -5.6387e-03, -2.0260e-03, -1.7095e-03,\n",
      "         8.4659e-04, -6.0866e-02, -5.6222e-02, -2.5379e-03, -1.5011e-02,\n",
      "        -1.2298e-03, -9.2632e-03,  3.1719e-03,  2.8508e-03,  4.7205e-03,\n",
      "         1.8103e-03, -4.7991e-03, -4.7207e-03, -2.6367e-03, -5.7746e-02,\n",
      "        -5.7507e-02, -5.5645e-02, -1.7834e-02, -3.9811e-03,  1.8550e-03,\n",
      "        -1.5613e-03, -5.8469e-02, -5.6094e-02, -2.0653e-02, -5.7249e-02,\n",
      "         2.4900e-03, -4.5148e-02, -4.9064e-02, -9.0694e-03,  1.7996e-02,\n",
      "        -5.2546e-02, -5.7078e-02, -1.6836e-02,  4.4030e-03, -5.0037e-02,\n",
      "         5.4468e-04, -4.6351e-02], device='cuda:0')), ('_blocklist.1._temporal_attention._U2', tensor([[-0.0181, -0.0188, -0.0193,  ..., -0.0183, -0.0198, -0.0197],\n",
      "        [-0.0374, -0.0390, -0.0657,  ..., -0.0461, -0.0613, -0.0475],\n",
      "        [-0.0124, -0.0127, -0.0123,  ..., -0.0120, -0.0124, -0.0126],\n",
      "        ...,\n",
      "        [ 0.0376,  0.0384,  0.0429,  ...,  0.0404,  0.0426,  0.0414],\n",
      "        [-0.0783, -0.0805, -0.0793,  ..., -0.0806, -0.0805, -0.0815],\n",
      "        [-0.0184, -0.0191, -0.0215,  ..., -0.0197, -0.0220, -0.0211]],\n",
      "       device='cuda:0')), ('_blocklist.1._temporal_attention._U3', tensor([ 1.5680e-01, -7.5301e-03,  3.7745e-02,  4.1136e-02, -2.3842e-02,\n",
      "         5.2928e-03, -1.7554e-03,  4.3219e-02,  2.0720e-04,  6.7793e-03,\n",
      "         2.5989e-02,  2.6644e-02,  1.6121e-01,  8.5170e-02,  1.9159e-01,\n",
      "         2.1699e-02,  4.6080e-02,  1.8933e-01, -1.4860e-02,  5.9868e-02,\n",
      "        -1.3018e-03,  1.0383e-01,  5.2814e-02,  3.5021e-02, -5.1293e-03,\n",
      "         9.0817e-03, -1.6799e-02, -6.1800e-03, -1.9222e-04,  2.5722e-02,\n",
      "         1.3500e-02,  3.9189e-02,  3.0658e-02,  1.7019e-01,  4.3929e-02,\n",
      "         8.5291e-02,  1.8348e-01, -8.0772e-03,  4.1137e-03, -1.0980e-02,\n",
      "         6.6017e-02,  6.6940e-03,  1.4360e-04,  4.9239e-03,  3.6300e-03,\n",
      "         1.2485e-02,  8.9307e-02,  6.9876e-03,  6.9557e-02,  2.8890e-01,\n",
      "        -4.2728e-03,  1.3462e-02, -3.4284e-03, -2.8527e-02,  2.5579e-01,\n",
      "         1.3130e-01,  4.0586e-02, -1.5876e-02,  4.9476e-02,  7.4618e-02,\n",
      "        -1.4011e-03,  5.1181e-03, -7.1166e-03,  5.3215e-02], device='cuda:0')), ('_blocklist.1._temporal_attention._be', tensor([[[-0.6767, -0.1888, -0.1722, -0.1581,  0.1407, -1.2172, -0.0496,\n",
      "           1.4203, -1.2246,  0.6632, -1.4235,  0.7341],\n",
      "         [ 0.2790,  0.3643,  0.7051,  0.1180,  0.4981,  0.0143, -0.0694,\n",
      "           0.0230,  0.8321, -1.8956, -0.4740, -0.7991],\n",
      "         [ 0.8475,  0.5805,  0.5964,  0.9006,  0.5540,  0.5712, -0.2636,\n",
      "          -1.5671,  0.8882, -0.6939, -0.2337, -0.9908],\n",
      "         [ 0.0054,  1.3601,  1.0587,  1.1426,  0.4834,  0.2976, -0.1903,\n",
      "           1.9385,  1.8745,  1.2903, -0.4386, -0.1874],\n",
      "         [-0.7356, -0.3502, -0.3158, -0.7967,  0.0031, -0.2750, -0.7841,\n",
      "          -0.6146,  0.5111, -0.7748,  1.0901,  0.5720],\n",
      "         [ 0.4119, -0.1818,  0.5317, -0.8954,  0.2689,  0.6035,  0.1914,\n",
      "          -1.6291,  1.2339,  1.4472,  1.6004, -1.0176],\n",
      "         [ 1.1768,  0.5733,  0.0765,  0.3578,  1.1063,  1.4411,  1.1324,\n",
      "           0.3371, -1.4919,  0.5587,  1.4758, -1.4864],\n",
      "         [ 0.0757, -1.4395, -0.9751, -1.2807, -0.3660,  0.4287,  0.8063,\n",
      "          -0.7718, -0.4672, -1.1803,  0.2369,  0.9979],\n",
      "         [-0.5643,  2.2152,  1.3466,  1.8785, -0.9459, -1.3815, -0.3993,\n",
      "           1.1871,  1.3388,  0.7483, -1.6945,  0.2185],\n",
      "         [-0.7394,  1.4929,  1.0187,  1.6734, -0.3324, -0.5128, -1.3544,\n",
      "           0.9155,  1.6804,  0.9393, -0.6059, -0.7071],\n",
      "         [ 1.1544, -0.4562, -0.5384, -0.7585,  1.2272,  1.0199,  1.3836,\n",
      "          -0.8166, -1.2780, -1.4385,  1.4995, -0.3862],\n",
      "         [ 0.7675, -0.5303, -0.2451, -0.6817,  1.5764,  1.6649,  1.3408,\n",
      "          -1.2483, -1.5217, -1.0381,  0.5377, -1.4916]]], device='cuda:0')), ('_blocklist.1._temporal_attention._Ve', tensor([[-0.4350, -0.2569, -0.5558, -0.5277, -0.1005, -0.4554, -0.6219, -0.4870,\n",
      "         -0.1525, -0.4263, -0.2871, -0.8118],\n",
      "        [ 0.1098,  0.2029,  0.2043, -0.1438,  0.2364, -0.0024, -0.3488,  0.3832,\n",
      "         -0.0648,  0.0971, -0.0298, -0.7639],\n",
      "        [ 0.1919, -0.0403,  0.1432,  0.2483, -0.0410, -0.3082,  0.0044,  0.1399,\n",
      "          0.1717,  0.4830, -0.3795, -0.3181],\n",
      "        [ 0.1545,  0.1428, -0.1273, -0.2616,  0.5228,  0.3921, -0.2380,  0.2103,\n",
      "          0.2729, -0.1636, -0.0138, -0.3702],\n",
      "        [ 0.2421,  0.0018, -0.1671,  0.4285,  0.6638, -0.1611, -0.3190, -0.1895,\n",
      "          0.0130,  0.2746,  0.0825, -0.2558],\n",
      "        [-0.2755, -0.1354, -0.0296,  0.3913, -0.0398,  0.4197,  0.0348,  0.2124,\n",
      "          0.0037,  0.2602, -0.1307,  0.0325],\n",
      "        [ 0.1775, -0.1417, -0.2974,  0.2613, -0.3733, -0.1450,  0.7287,  0.2933,\n",
      "          0.4527, -0.1690,  0.1409,  0.0260],\n",
      "        [-0.1273, -0.0407,  0.4303, -0.1286,  0.1451,  0.1763, -0.1893,  0.6471,\n",
      "          0.2412,  0.1596,  0.0496, -0.1273],\n",
      "        [ 0.5450,  0.0372,  0.0534,  0.1290, -0.0508, -0.1645,  0.5163, -0.1074,\n",
      "         -0.3939,  0.2396,  0.2626,  0.2085],\n",
      "        [ 0.0350,  0.4350,  0.3041,  0.0482,  0.2679, -0.2539,  0.0457,  0.0111,\n",
      "          0.4515, -0.4038,  0.3577,  0.0571],\n",
      "        [-0.1013, -0.0669, -0.0649,  0.1857,  0.3663,  0.2015,  0.3184, -0.2837,\n",
      "          0.0682,  0.2896,  0.8526, -0.1683],\n",
      "        [-0.1878, -0.0021,  0.1557,  0.2665, -0.1115,  0.0251,  0.2797, -0.3720,\n",
      "          0.3616,  0.4333,  0.0860,  0.5453]], device='cuda:0')), ('_blocklist.1._spatial_attention._W1', tensor([-0.0267,  0.1914,  0.0183,  0.1050, -0.1100, -0.1226, -0.1992,  0.1210,\n",
      "         0.0421,  0.1897, -0.0996,  0.0190], device='cuda:0')), ('_blocklist.1._spatial_attention._W2', tensor([[-1.5178e-01,  2.1654e-01, -1.6608e-02,  2.5569e-02, -1.8865e-01,\n",
      "         -8.3805e-02,  3.1986e-02,  1.8469e-01, -6.0479e-02,  1.5712e-01,\n",
      "         -4.6157e-02, -2.3845e-02],\n",
      "        [ 2.0633e-02, -2.0859e-01,  1.1447e-01,  2.2316e-01,  4.7879e-02,\n",
      "          8.5434e-02,  2.5634e-01,  7.5948e-02, -2.1758e-01,  1.9277e-01,\n",
      "          1.4576e-01, -1.7914e-01],\n",
      "        [ 6.3119e-02, -2.3708e-01,  1.0435e-02, -5.9091e-02, -1.6165e-01,\n",
      "          1.0413e-01, -2.6600e-02, -1.2810e-01, -6.6634e-02, -6.9477e-02,\n",
      "          2.6875e-02, -2.1496e-02],\n",
      "        [-3.5059e-02, -6.6618e-02, -3.9525e-02,  9.8826e-03, -1.3823e-01,\n",
      "          1.0775e-01,  2.1111e-01, -5.7535e-02, -1.2615e-03, -6.9927e-02,\n",
      "          2.7274e-02, -2.3259e-01],\n",
      "        [-1.3778e-01,  1.9517e-01, -1.4771e-01,  1.3162e-02,  4.2899e-02,\n",
      "         -9.2816e-02, -2.1353e-01,  2.2909e-01, -3.4644e-02,  1.4746e-01,\n",
      "         -4.2803e-02,  1.6159e-02],\n",
      "        [ 8.3576e-02,  7.3185e-02,  9.6170e-02,  1.6239e-03,  1.0492e-03,\n",
      "          1.3603e-01,  5.2161e-02,  1.2525e-01, -2.1386e-02, -6.3541e-03,\n",
      "          5.4157e-02,  1.1185e-01],\n",
      "        [-1.9930e-01, -2.2916e-01, -2.6506e-01, -3.6532e-01, -1.3112e-01,\n",
      "         -4.2971e-02, -3.4733e-01, -1.5147e-01, -1.4988e-01, -2.1103e-01,\n",
      "         -3.1052e-01, -3.1970e-01],\n",
      "        [-3.2446e-03,  1.0106e-01,  5.9525e-02, -4.5256e-02,  1.8005e-01,\n",
      "         -2.2892e-01, -1.8270e-01,  1.6261e-02,  1.8555e-01, -8.8848e-02,\n",
      "         -3.1805e-02,  1.9231e-01],\n",
      "        [-1.9844e-01, -7.1560e-03, -1.0963e-01,  7.8808e-02, -9.9101e-02,\n",
      "         -7.9996e-02,  4.4429e-02,  5.7584e-02, -3.5583e-02,  7.7186e-02,\n",
      "          3.0217e-02,  9.7850e-03],\n",
      "        [ 1.0717e-01,  1.5705e-01, -7.0094e-03,  2.5524e-02, -5.7837e-03,\n",
      "         -2.4361e-02, -1.2901e-03, -1.2005e-01, -1.7994e-01,  2.9410e-02,\n",
      "          2.0638e-02, -1.7357e-01],\n",
      "        [ 2.7631e-02, -2.0072e-01,  2.0350e-01, -1.2062e-01,  2.6521e-02,\n",
      "          3.2526e-02, -6.5808e-02, -1.9351e-01, -1.9425e-01,  2.1453e-02,\n",
      "          1.2164e-01, -6.0459e-03],\n",
      "        [ 4.7435e-02,  1.9140e-01, -6.2222e-03, -1.9025e-01, -5.7139e-02,\n",
      "          8.9958e-03,  5.6856e-03, -1.1155e-01, -4.2235e-02, -2.1983e-01,\n",
      "          1.0360e-01, -1.3404e-01],\n",
      "        [ 4.0638e-02,  1.6703e-01,  8.1462e-02,  4.2291e-03,  8.8442e-02,\n",
      "         -7.0142e-02, -2.7811e-01, -6.9757e-02, -1.1411e-01, -2.8844e-01,\n",
      "         -1.0049e-01, -1.8812e-01],\n",
      "        [ 7.4899e-03, -1.8688e-02,  1.6846e-01, -5.8899e-03,  3.8123e-02,\n",
      "          2.0821e-01,  6.4553e-02, -1.8074e-01, -1.4767e-01,  1.3794e-01,\n",
      "          1.5662e-02,  1.8013e-01],\n",
      "        [ 3.4389e-01,  2.7298e-01,  2.8594e-01,  1.1631e-01,  3.2950e-01,\n",
      "          3.2472e-01, -1.2280e-01,  2.7479e-01,  1.3351e-01,  1.1563e-01,\n",
      "          1.1654e-01, -6.9439e-02],\n",
      "        [ 3.4103e-02,  1.7397e-01, -2.8305e-01, -6.1077e-02, -3.9453e-02,\n",
      "         -1.5948e-01, -6.9239e-02,  6.7552e-03,  1.9178e-01, -2.7149e-01,\n",
      "          1.3535e-01, -1.1685e-01],\n",
      "        [-5.0919e-02, -2.1382e-01, -2.4176e-01, -1.4661e-01, -4.4244e-02,\n",
      "          1.5131e-01, -2.6909e-01, -6.3729e-02, -4.4084e-02, -5.6208e-02,\n",
      "         -1.0635e-01, -2.0690e-01],\n",
      "        [ 2.4539e-02,  7.2917e-02,  9.6283e-02,  1.8720e-01, -3.6803e-02,\n",
      "         -1.2487e-03, -6.3176e-02, -8.0497e-03,  1.8356e-01,  2.3341e-02,\n",
      "         -2.0225e-01,  8.7755e-02],\n",
      "        [-1.7467e-01, -2.9347e-01, -5.7015e-03, -2.7822e-01, -1.7100e-01,\n",
      "         -1.7753e-01, -1.8883e-01, -3.9880e-01, -1.9009e-01, -1.8886e-01,\n",
      "         -2.6019e-01, -1.4354e-01],\n",
      "        [ 4.7921e-02, -2.6750e-02, -2.9987e-02, -2.3528e-01,  1.1672e-02,\n",
      "          1.3194e-01, -1.3895e-02,  1.8142e-02,  3.4310e-02,  9.9915e-02,\n",
      "         -1.0883e-01, -3.4374e-02],\n",
      "        [-1.1907e-02,  1.0349e-01, -1.8761e-01, -9.6659e-02,  1.9905e-01,\n",
      "         -1.1771e-01,  1.5538e-01, -1.4746e-01, -2.2786e-02, -1.5256e-01,\n",
      "         -4.7115e-02,  2.8554e-02],\n",
      "        [-1.7047e-01,  1.6562e-01,  2.3605e-01,  2.3510e-01,  2.3690e-01,\n",
      "          5.1096e-02,  2.1803e-01, -1.4227e-01, -1.0862e-02, -5.1276e-03,\n",
      "         -4.8510e-02,  1.9204e-02],\n",
      "        [ 8.1249e-02,  1.5146e-02,  1.6906e-01, -5.0488e-02, -2.5651e-01,\n",
      "         -3.9781e-02, -3.8323e-02, -1.4728e-01, -1.8205e-01, -2.8497e-02,\n",
      "         -1.2576e-01, -1.6347e-01],\n",
      "        [-1.4059e-01, -2.2071e-01, -4.2955e-02,  4.2932e-02, -4.0415e-02,\n",
      "         -5.0161e-02,  1.3023e-02, -6.3583e-02,  1.3070e-01,  4.9320e-02,\n",
      "          4.5485e-02,  2.7110e-02],\n",
      "        [ 2.3905e-02,  1.4982e-01, -1.6257e-01,  4.7662e-02,  1.8913e-01,\n",
      "          6.0273e-02, -1.5138e-01,  2.7704e-01, -4.1294e-03,  2.0388e-01,\n",
      "          1.7626e-01, -1.8794e-01],\n",
      "        [ 1.6560e-01, -6.5262e-02,  2.3027e-01, -1.9830e-01, -1.6974e-01,\n",
      "          1.4701e-01,  1.9300e-01, -9.1208e-02, -8.3982e-03, -1.8023e-01,\n",
      "          5.4751e-02, -7.9372e-02],\n",
      "        [-9.1253e-02,  1.0701e-01,  2.4295e-02, -1.3444e-01,  1.8880e-01,\n",
      "         -1.6520e-01, -6.3208e-02,  1.7977e-01,  1.8179e-01, -1.7062e-01,\n",
      "          2.8362e-01, -6.2710e-02],\n",
      "        [ 6.6937e-02, -2.1190e-01,  1.0588e-01,  5.0004e-02,  1.4894e-02,\n",
      "          4.6708e-02, -9.8077e-03,  1.2778e-01, -2.0777e-02,  2.0302e-02,\n",
      "         -4.4398e-02,  3.5147e-02],\n",
      "        [-1.1157e-02, -8.1877e-02,  5.8073e-02, -2.7280e-01,  3.1255e-02,\n",
      "         -1.7970e-01, -1.4527e-01, -1.9933e-01, -1.7109e-01, -3.7048e-01,\n",
      "         -2.9981e-02, -5.7634e-02],\n",
      "        [ 3.1975e-02, -3.0040e-02,  1.2752e-01, -7.7097e-02, -7.9525e-02,\n",
      "          1.3158e-01,  1.8790e-01, -3.1712e-02, -1.5598e-01,  1.4296e-01,\n",
      "         -8.0250e-02, -2.4378e-01],\n",
      "        [ 3.8810e-03, -1.7366e-01,  1.3212e-01,  2.4059e-01,  1.4802e-02,\n",
      "          2.0063e-02, -1.6457e-01,  3.0317e-02, -2.4537e-02,  1.8151e-01,\n",
      "         -1.0036e-01, -1.0469e-01],\n",
      "        [-1.4688e-01, -2.4596e-01, -2.2753e-01, -1.3004e-01, -7.4319e-02,\n",
      "         -3.2829e-01, -2.5477e-01,  1.3930e-02,  7.4374e-02, -2.4008e-01,\n",
      "         -2.3612e-01,  9.0090e-02],\n",
      "        [-2.1292e-01,  1.1606e-01,  1.7429e-01,  3.1699e-02,  9.0096e-02,\n",
      "          1.1129e-01,  1.6452e-01, -1.0637e-01, -1.5608e-01,  6.8736e-02,\n",
      "         -1.4495e-02,  2.0402e-01],\n",
      "        [-1.1239e-01,  2.5619e-02,  4.0649e-02,  2.7220e-02, -3.2650e-02,\n",
      "         -8.8255e-02,  8.6565e-02,  2.1337e-01,  2.3244e-02, -7.8336e-02,\n",
      "          3.5283e-02, -5.2641e-02],\n",
      "        [-1.5885e-01, -1.8543e-01, -1.6566e-01, -1.4622e-01, -3.5781e-02,\n",
      "         -1.4695e-01, -1.3919e-01, -1.6521e-01,  1.0459e-03, -2.2513e-01,\n",
      "         -1.3855e-01, -3.0590e-01],\n",
      "        [-7.2654e-02,  3.5648e-02,  2.0225e-01, -2.2925e-01, -3.0162e-02,\n",
      "          5.9716e-02,  7.4390e-02,  4.6247e-02,  3.7357e-02,  1.0415e-02,\n",
      "          4.8020e-02,  6.3714e-02],\n",
      "        [ 7.6062e-03,  1.1806e-03,  9.1194e-02, -1.6286e-01, -1.7485e-01,\n",
      "         -9.4578e-02, -7.2187e-02,  1.6367e-01,  6.0679e-02,  2.2747e-01,\n",
      "          6.6055e-03, -2.3291e-01],\n",
      "        [ 7.8786e-02,  2.6185e-01, -1.1203e-01,  2.7147e-01, -3.7275e-02,\n",
      "         -5.2867e-02,  1.3553e-01,  1.4332e-01, -3.2457e-02,  8.8781e-02,\n",
      "         -2.0737e-02,  1.8984e-01],\n",
      "        [ 4.4417e-01,  3.6822e-01,  8.5476e-02,  4.0968e-01, -9.2728e-03,\n",
      "          2.9934e-01,  6.9360e-02,  4.0844e-01,  2.1939e-01,  2.9459e-01,\n",
      "          2.5981e-02,  2.8836e-01],\n",
      "        [-1.3350e-01, -3.8117e-02,  2.3133e-01,  1.2801e-01, -1.8805e-01,\n",
      "          1.2673e-01,  2.1246e-02, -2.0420e-01, -8.0977e-02,  9.9217e-02,\n",
      "          3.3726e-02,  2.1636e-01],\n",
      "        [-7.0159e-02, -2.9232e-02, -3.2279e-02, -2.1177e-01,  4.0539e-04,\n",
      "         -1.4856e-01, -8.2390e-02, -2.2057e-01,  1.8789e-02,  1.7365e-01,\n",
      "         -5.0647e-03, -1.8953e-01],\n",
      "        [-4.7179e-02, -1.8727e-01, -8.2136e-02, -1.3252e-01,  1.7692e-01,\n",
      "          1.7615e-01, -9.4660e-02, -5.8864e-04, -1.7084e-01, -1.9755e-01,\n",
      "         -5.1713e-02, -9.8318e-02],\n",
      "        [ 5.0600e-03, -2.4288e-03, -1.4646e-03,  9.7321e-02, -7.3344e-04,\n",
      "          1.7427e-01, -1.9564e-01, -1.1460e-01,  6.5015e-02, -1.9198e-01,\n",
      "          1.5962e-01,  1.1769e-01],\n",
      "        [-1.2502e-01,  2.1631e-02,  1.2231e-01, -1.1100e-01, -2.0585e-01,\n",
      "         -5.2127e-03,  1.5958e-02, -2.0939e-01, -1.9190e-01, -8.9840e-04,\n",
      "          1.8252e-01, -2.1813e-01],\n",
      "        [ 9.4709e-02, -1.4813e-01, -1.3542e-01, -1.3450e-02,  1.6593e-03,\n",
      "          1.2218e-01,  8.1253e-02,  3.1374e-02, -1.8002e-02, -2.7564e-03,\n",
      "          7.6498e-02,  2.0788e-01],\n",
      "        [ 4.6217e-02, -1.1618e-01, -9.7327e-02, -7.1609e-02, -1.4590e-01,\n",
      "         -1.2416e-01, -2.3118e-01, -1.7775e-03,  1.4783e-01,  2.4340e-02,\n",
      "         -6.2298e-03,  2.3206e-01],\n",
      "        [ 1.9303e-01,  6.7212e-02, -4.0902e-03, -1.2750e-01, -1.0166e-01,\n",
      "         -1.8423e-01, -1.2180e-03,  1.6919e-01,  2.9502e-02,  2.2948e-01,\n",
      "         -2.2199e-01, -9.7095e-03],\n",
      "        [-4.8420e-03, -1.7966e-01, -1.9544e-01, -1.6463e-01, -2.3400e-01,\n",
      "         -1.9342e-01, -8.7482e-03, -1.7164e-02, -1.4953e-01, -1.2145e-01,\n",
      "          2.1677e-01, -1.3744e-01],\n",
      "        [-1.1266e-01,  1.9137e-01, -4.3344e-02, -2.2413e-01,  1.0024e-01,\n",
      "          1.9055e-01, -2.2573e-01,  5.6671e-02,  4.4940e-02, -2.5908e-02,\n",
      "         -8.7358e-02,  2.1452e-01],\n",
      "        [-7.4460e-02, -1.5313e-01, -2.1241e-01,  9.3650e-02, -1.8663e-01,\n",
      "         -2.0554e-01, -4.6375e-02, -1.5558e-01, -2.8434e-01, -2.0392e-01,\n",
      "          1.1945e-01,  6.6767e-02],\n",
      "        [-1.0348e-02, -6.5573e-03, -2.4034e-01, -6.3146e-02, -1.2327e-01,\n",
      "         -1.0203e-01,  1.5326e-01, -2.1831e-01,  1.5971e-01, -1.0457e-01,\n",
      "          9.0758e-02, -1.1231e-02],\n",
      "        [ 1.3620e-02,  1.2570e-01,  9.8645e-02, -8.3000e-02, -2.2724e-01,\n",
      "          1.0705e-01, -2.6537e-03, -8.9511e-02, -1.5243e-01, -2.1286e-01,\n",
      "          3.6319e-02,  9.9806e-02],\n",
      "        [ 1.9914e-01, -2.4354e-01, -6.1301e-02,  3.7128e-02, -1.9849e-01,\n",
      "         -2.3687e-01, -1.4306e-01,  5.2590e-02,  4.7258e-02, -2.1865e-01,\n",
      "          2.7186e-02, -8.2814e-02],\n",
      "        [-4.9207e-03, -1.8095e-02,  3.3084e-02,  1.6462e-01,  4.7168e-02,\n",
      "         -1.4402e-01,  1.6976e-01, -1.4205e-02, -1.4876e-01, -7.8575e-02,\n",
      "         -2.2011e-01, -2.1321e-01],\n",
      "        [ 7.6712e-02, -1.7122e-02,  4.7986e-02, -3.6524e-02, -4.4546e-03,\n",
      "         -1.6703e-01, -8.8704e-03, -2.7500e-02, -2.5894e-01, -1.5371e-01,\n",
      "         -3.2789e-02, -1.5423e-01],\n",
      "        [ 8.3111e-02,  1.9444e-01,  1.6709e-01, -2.1430e-01, -1.2643e-01,\n",
      "          4.6133e-03, -1.3675e-01, -1.8665e-01,  1.0512e-01,  2.4437e-01,\n",
      "         -1.7869e-01,  6.0945e-02],\n",
      "        [-1.0501e-01,  2.7478e-02, -4.4058e-02,  7.0467e-02, -1.1239e-01,\n",
      "         -4.7667e-03,  4.2766e-03, -1.9585e-01,  1.6060e-01,  5.8414e-02,\n",
      "         -1.3735e-01, -2.0931e-02],\n",
      "        [ 2.1983e-01,  2.4204e-01,  1.2591e-01,  6.0058e-02,  1.5657e-01,\n",
      "         -1.2421e-01, -5.4354e-02,  9.0417e-02,  2.2278e-01, -1.3650e-01,\n",
      "          7.1290e-02,  2.7886e-02],\n",
      "        [ 6.8343e-02,  2.2968e-01, -1.8984e-01,  2.2217e-01, -1.6682e-01,\n",
      "         -1.0181e-01, -1.6994e-03,  3.1686e-02,  1.0865e-01, -7.2687e-03,\n",
      "         -6.2613e-04, -1.6430e-02],\n",
      "        [ 1.3093e-01,  1.6595e-01, -6.2292e-02, -1.8038e-01,  1.1401e-01,\n",
      "          4.9591e-03, -2.2540e-02,  2.7656e-02,  1.5379e-02, -3.2000e-02,\n",
      "          7.7780e-02,  1.9845e-01],\n",
      "        [-4.5602e-02, -5.1029e-02,  4.7157e-02,  2.0289e-01,  1.1524e-01,\n",
      "          1.2985e-01,  3.1793e-01, -5.2830e-02, -5.6321e-02,  8.4582e-02,\n",
      "          1.8849e-01, -1.3189e-02],\n",
      "        [ 1.1351e-01,  5.5052e-02, -3.9878e-02,  1.6961e-01,  5.1142e-02,\n",
      "         -2.4845e-01, -1.6271e-02,  4.6343e-02,  1.6087e-01,  1.3859e-01,\n",
      "          1.5986e-01, -1.7527e-01],\n",
      "        [-7.9504e-02,  1.8945e-01, -9.4780e-02, -1.2593e-01, -6.1410e-03,\n",
      "          1.5232e-01,  1.8653e-02,  2.1587e-02,  1.5689e-01, -1.2440e-01,\n",
      "          2.3905e-01,  1.5299e-01],\n",
      "        [ 3.8408e-02, -5.5261e-02, -1.4507e-01,  1.4305e-02,  2.0865e-01,\n",
      "          1.2466e-01, -6.4449e-02, -2.4910e-01, -6.9399e-02,  1.4563e-01,\n",
      "         -2.8818e-02,  9.8703e-02]], device='cuda:0')), ('_blocklist.1._spatial_attention._W3', tensor([-0.0335,  0.4712,  0.1694,  0.2092,  0.2511, -0.2567, -0.3622, -0.1679,\n",
      "        -0.1686,  0.3766, -0.0246, -0.2079, -0.1030,  0.0164,  0.2178,  0.5038,\n",
      "        -0.2425,  0.5134, -0.3990, -0.2549,  0.4553, -0.1571,  0.1789,  0.2203,\n",
      "         0.9237,  0.1981,  0.1457,  0.3075, -0.3342,  0.2070, -0.0374, -0.2908,\n",
      "        -0.2415,  0.2787, -0.2680,  0.6511, -0.1615,  0.6788,  0.2028,  0.0033,\n",
      "         0.3579,  0.7493,  0.1203,  0.4900, -0.2199,  0.8376,  0.2455,  0.6190,\n",
      "         0.1849,  0.2907,  0.5041,  0.2960,  0.2234,  0.4259, -0.0805,  0.0657,\n",
      "         0.2211,  0.5894,  0.4550,  0.2335,  0.2335,  0.6480,  0.3911,  0.1627],\n",
      "       device='cuda:0')), ('_blocklist.1._spatial_attention._bs', tensor([[[ 0.4040, -0.4125, -0.0593,  ...,  0.3821, -0.2496, -0.7324],\n",
      "         [-0.4481,  0.0807, -0.4843,  ..., -0.2136,  0.0022, -0.0329],\n",
      "         [ 0.0434, -0.0407, -0.3908,  ..., -0.3513,  0.6205,  0.6184],\n",
      "         ...,\n",
      "         [-0.3098, -0.2402, -0.5437,  ...,  0.0233, -0.0700,  0.1013],\n",
      "         [ 0.3617,  0.2412, -0.3462,  ...,  0.2231,  0.4608, -0.2235],\n",
      "         [-0.2773, -0.3593,  0.0852,  ..., -0.1275, -0.2296,  0.1290]]],\n",
      "       device='cuda:0')), ('_blocklist.1._spatial_attention._Vs', tensor([[ 0.2307,  0.0968, -0.0971,  ...,  0.0947, -0.0272,  0.0763],\n",
      "        [ 0.0626,  0.4923, -0.0586,  ...,  0.1123, -0.0699, -0.0671],\n",
      "        [-0.0025,  0.0317,  0.2070,  ...,  0.1219,  0.1908,  0.0171],\n",
      "        ...,\n",
      "        [ 0.1541,  0.0468, -0.0637,  ...,  0.2147,  0.0663, -0.0696],\n",
      "        [ 0.0148,  0.1050,  0.0146,  ..., -0.0377,  0.2768,  0.0092],\n",
      "        [ 0.0432,  0.1340,  0.0603,  ..., -0.0018,  0.0280,  0.4929]],\n",
      "       device='cuda:0')), ('_blocklist.1._chebconv_attention._weight', tensor([[[ 7.5286e-01, -1.3260e+00,  3.7664e-01,  ...,  1.3694e-01,\n",
      "          -9.2797e-01, -1.4515e-01],\n",
      "         [ 6.7757e-01, -1.2624e+00,  3.9623e-01,  ...,  1.4986e-01,\n",
      "          -9.1418e-01, -1.6755e-01],\n",
      "         [ 7.3239e-01, -1.1552e+00,  4.0860e-01,  ...,  2.0919e-02,\n",
      "          -8.8002e-01, -1.0452e-01],\n",
      "         ...,\n",
      "         [-4.7265e-01,  1.3682e+00, -4.7251e-01,  ..., -1.1165e-01,\n",
      "           7.1899e-01,  1.3194e-01],\n",
      "         [ 7.0723e-01, -1.3588e+00,  3.8817e-01,  ...,  1.9774e-01,\n",
      "          -9.6823e-01, -2.4091e-01],\n",
      "         [ 6.6140e-01, -1.2462e+00,  2.5008e-01,  ...,  9.2695e-02,\n",
      "          -7.3535e-01, -1.1290e-01]],\n",
      "\n",
      "        [[-6.0076e-01,  2.1445e-01,  7.7733e-01,  ...,  2.0957e+00,\n",
      "           4.8871e-02,  1.8668e-01],\n",
      "         [-2.1663e-01,  2.8360e-02,  3.0985e-02,  ...,  2.3866e-01,\n",
      "           5.3300e-03,  1.1724e-02],\n",
      "         [-5.6543e-01,  2.4458e-01,  1.0498e+00,  ...,  2.3179e+00,\n",
      "           3.6644e-02,  2.7146e-01],\n",
      "         ...,\n",
      "         [ 3.3604e-01, -4.2234e-02, -8.9253e-02,  ..., -5.2868e-01,\n",
      "          -4.7776e-03, -3.8235e-02],\n",
      "         [-5.0342e-01,  1.1464e-01,  1.8110e-01,  ...,  1.0947e+00,\n",
      "           5.1112e-02,  7.9840e-02],\n",
      "         [-6.4470e-02,  1.0370e-02,  4.6035e-02,  ...,  1.0745e-01,\n",
      "          -3.8745e-04,  3.3550e-03]],\n",
      "\n",
      "        [[-5.8605e-01,  1.8490e+00, -5.5217e-01,  ..., -1.5008e-01,\n",
      "           7.3695e-01,  1.1046e-01],\n",
      "         [-4.8383e-01,  1.6479e+00, -4.6574e-01,  ..., -1.8235e-01,\n",
      "           7.0746e-01,  1.2998e-01],\n",
      "         [-5.9035e-01,  1.5246e+00, -4.9803e-01,  ..., -7.6015e-02,\n",
      "           6.2701e-01,  2.4298e-02],\n",
      "         ...,\n",
      "         [ 5.0237e-01, -1.7019e+00,  5.3919e-01,  ...,  2.0067e-01,\n",
      "          -6.6787e-01, -1.2215e-01],\n",
      "         [-5.1246e-01,  1.7989e+00, -4.8025e-01,  ..., -2.1901e-01,\n",
      "           7.9104e-01,  1.5744e-01],\n",
      "         [-5.0774e-01,  1.5907e+00, -4.1479e-01,  ..., -1.0669e-01,\n",
      "           5.7658e-01,  5.6804e-02]]], device='cuda:0')), ('_blocklist.1._chebconv_attention._bias', tensor([ 9.2641e-02,  2.2257e-01,  1.7075e-01,  2.2832e-01,  1.3845e-01,\n",
      "         8.0220e-01,  2.0259e-01,  1.1788e-01,  2.5269e-01,  2.8223e-01,\n",
      "         5.4295e-01,  9.3106e-02,  2.3373e-01,  1.6853e-01,  1.2677e-01,\n",
      "         3.5840e-01,  1.3584e-01,  6.3703e-01,  5.9869e-02,  6.3499e-01,\n",
      "         3.2842e-02,  6.3447e-01,  3.8919e-01,  6.4924e-01,  4.9239e-01,\n",
      "         3.9277e-01,  4.8297e-01,  1.0373e-01,  3.1487e-02,  2.5509e-01,\n",
      "         8.4009e-02,  5.0176e-01,  3.2478e-01,  4.6936e-01,  1.2158e-01,\n",
      "         3.1994e-01,  3.2155e-02,  5.2958e-01,  5.2458e-01,  2.3254e-01,\n",
      "         2.9128e-01,  2.9819e-01,  1.9437e-01,  5.9857e-01,  3.2782e-01,\n",
      "         1.9191e-01,  1.8539e-01,  3.3479e-01,  7.5984e-01,  3.2184e-01,\n",
      "         6.0605e-01,  2.6781e-01,  2.0565e-01,  2.0230e-01,  1.0949e-01,\n",
      "         3.2641e-01,  4.3302e-01, -4.8857e-04,  3.1879e-01,  9.6089e-02,\n",
      "         2.1326e-01,  5.1881e-01,  1.5493e-01,  4.9666e-01], device='cuda:0')), ('_blocklist.1._time_convolution.weight', tensor([[[[-0.4070, -0.0437, -0.3790]],\n",
      "\n",
      "         [[ 0.1845, -0.0697,  0.2560]],\n",
      "\n",
      "         [[-0.1921, -0.0973, -0.2550]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-0.0938, -0.0522,  0.0175]],\n",
      "\n",
      "         [[ 0.2693,  0.0629,  0.1646]],\n",
      "\n",
      "         [[ 0.0726, -0.0722, -0.1210]]],\n",
      "\n",
      "\n",
      "        [[[-0.0310, -0.1417,  0.0609]],\n",
      "\n",
      "         [[-0.1866, -0.2158, -0.0550]],\n",
      "\n",
      "         [[-0.0622, -0.0911,  0.0648]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.0637,  0.1181,  0.1027]],\n",
      "\n",
      "         [[-0.0496, -0.1014, -0.0629]],\n",
      "\n",
      "         [[ 0.1409,  0.0240,  0.0778]]],\n",
      "\n",
      "\n",
      "        [[[-0.2843, -0.1550, -0.1959]],\n",
      "\n",
      "         [[-0.0386,  0.2709,  0.0194]],\n",
      "\n",
      "         [[-0.0791,  0.0244, -0.0199]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.0896, -0.0763, -0.0197]],\n",
      "\n",
      "         [[ 0.1133,  0.3445,  0.1699]],\n",
      "\n",
      "         [[-0.0029, -0.0623, -0.1007]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 0.1917,  0.2745,  0.3254]],\n",
      "\n",
      "         [[-0.2098, -0.2522, -0.4066]],\n",
      "\n",
      "         [[ 0.1700,  0.1274,  0.0342]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.1381,  0.0032, -0.1205]],\n",
      "\n",
      "         [[-0.3433, -0.3383, -0.3083]],\n",
      "\n",
      "         [[ 0.0635,  0.0571, -0.0410]]],\n",
      "\n",
      "\n",
      "        [[[-0.0586,  0.0103, -0.0676]],\n",
      "\n",
      "         [[-0.0208, -0.0602, -0.0659]],\n",
      "\n",
      "         [[ 0.0735, -0.0525,  0.0844]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-0.0908, -0.1165,  0.1246]],\n",
      "\n",
      "         [[-0.0400,  0.0091,  0.0505]],\n",
      "\n",
      "         [[-0.0248, -0.0079,  0.0018]]],\n",
      "\n",
      "\n",
      "        [[[-0.2441, -0.0343,  0.1278]],\n",
      "\n",
      "         [[ 0.1978, -0.5657, -0.4233]],\n",
      "\n",
      "         [[-0.0457,  0.0417,  0.1193]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.0666,  0.1459,  0.0376]],\n",
      "\n",
      "         [[ 0.2261, -0.3249, -0.5612]],\n",
      "\n",
      "         [[ 0.0347,  0.0433, -0.0472]]]], device='cuda:0')), ('_blocklist.1._time_convolution.bias', tensor([ 0.7959,  0.6304,  0.8504,  0.5612,  0.0475,  0.1203,  0.7759,  0.1201,\n",
      "         0.1238,  0.6339,  0.9838,  0.1909,  0.3645,  0.4177,  0.1481,  0.8099,\n",
      "         0.2179,  0.5221,  0.9855,  0.7790,  0.8117,  0.9198,  0.4655,  0.4128,\n",
      "         0.7791,  0.9589,  0.5969,  0.1206,  0.2000,  0.2352,  0.8479,  0.4410,\n",
      "        -0.0275,  0.4166,  0.8086,  0.4863,  0.4478,  0.5110,  0.6612,  0.4945,\n",
      "         0.2492,  0.9193,  0.8368,  0.6685,  0.6016,  0.0997,  0.7388,  0.3662,\n",
      "         0.1201,  0.3101,  0.9249,  0.6236,  0.9011,  0.3593,  0.8973,  0.4448,\n",
      "         0.8095,  0.3048,  0.0374, -0.0384,  0.7013,  0.1314,  0.3274,  0.5876],\n",
      "       device='cuda:0')), ('_blocklist.1._residual_convolution.weight', tensor([[[[-0.0642]],\n",
      "\n",
      "         [[ 0.1826]],\n",
      "\n",
      "         [[-0.9962]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.0923]],\n",
      "\n",
      "         [[ 0.0740]],\n",
      "\n",
      "         [[ 0.3614]]],\n",
      "\n",
      "\n",
      "        [[[ 0.1166]],\n",
      "\n",
      "         [[-0.0968]],\n",
      "\n",
      "         [[-0.1258]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-0.0975]],\n",
      "\n",
      "         [[-0.1563]],\n",
      "\n",
      "         [[-0.1745]]],\n",
      "\n",
      "\n",
      "        [[[ 0.2414]],\n",
      "\n",
      "         [[ 0.0477]],\n",
      "\n",
      "         [[ 0.1508]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-0.2601]],\n",
      "\n",
      "         [[-0.0230]],\n",
      "\n",
      "         [[-0.1665]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 0.0496]],\n",
      "\n",
      "         [[-0.1714]],\n",
      "\n",
      "         [[ 0.1799]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-0.1947]],\n",
      "\n",
      "         [[ 0.2222]],\n",
      "\n",
      "         [[-0.2535]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0231]],\n",
      "\n",
      "         [[ 0.1922]],\n",
      "\n",
      "         [[-0.1015]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-0.0507]],\n",
      "\n",
      "         [[ 0.2183]],\n",
      "\n",
      "         [[ 0.1052]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0730]],\n",
      "\n",
      "         [[ 0.1521]],\n",
      "\n",
      "         [[-0.4923]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.2136]],\n",
      "\n",
      "         [[-0.0511]],\n",
      "\n",
      "         [[ 0.1117]]]], device='cuda:0')), ('_blocklist.1._residual_convolution.bias', tensor([ 0.7786,  0.8710,  0.0436,  0.4098,  0.5208,  0.2829,  0.6250, -0.0321,\n",
      "         0.0255,  0.6474,  0.6357,  0.1798,  0.8722,  0.3000,  0.1078,  0.3516,\n",
      "         0.9122,  0.9175,  0.7630,  0.5078,  0.6716,  0.1770, -0.0192,  0.5175,\n",
      "         0.7000,  0.5158,  0.0393,  0.6300,  0.7098,  0.4411,  0.0309,  0.3441,\n",
      "        -0.0274,  0.9798,  1.0098,  0.5670,  0.9350,  0.9428,  0.2060,  0.5262,\n",
      "         0.3860,  0.0951,  0.9384,  0.7268,  0.0689,  0.1314,  0.7499,  0.6230,\n",
      "         0.8360,  0.3221,  0.2616,  0.4445,  0.1502,  0.7997,  0.6585,  0.6089,\n",
      "        -0.0104,  0.3658,  0.4143,  0.4371,  0.7024,  0.2402,  0.8692,  0.4904],\n",
      "       device='cuda:0')), ('_blocklist.1._layer_norm.weight', tensor([ 1.0964, -0.9414,  1.6469,  1.0567, -0.8049,  1.4955,  1.2203,  1.4977,\n",
      "         0.8479, -0.8264,  1.1292,  1.3754,  1.3494,  1.5497,  0.9967, -0.7920,\n",
      "         1.6761,  1.5365,  1.6924, -0.5933,  1.5730,  0.8949,  1.3725,  1.2712,\n",
      "         1.3994,  1.4541,  0.8749,  1.0203,  1.1294, -0.7011,  1.4323,  1.5222,\n",
      "         1.5032, -0.5448,  0.9804,  1.4091,  1.6352,  1.2381, -0.6033,  1.1310,\n",
      "         1.3112,  1.3462,  1.7739,  1.0304,  1.5302, -0.6275,  1.3582,  1.2260,\n",
      "         1.3359,  1.3877,  1.5075,  1.4143,  1.5840,  1.2077,  1.1437,  1.3788,\n",
      "         1.4825, -0.3774,  1.3252,  1.4717,  1.4294, -0.4181,  0.8878,  1.3706],\n",
      "       device='cuda:0')), ('_blocklist.1._layer_norm.bias', tensor([ 0.9339,  0.6209,  0.8357,  0.9937,  1.0148, -0.4533,  1.4809, -0.2454,\n",
      "        -0.2072,  0.6596,  1.4187,  1.2913,  0.8704,  1.2348,  0.9977,  1.0059,\n",
      "        -0.3151,  0.9441, -0.4363,  1.3752,  0.8674,  1.1402,  1.4172,  0.7437,\n",
      "        -0.5219,  0.9024, -0.5441,  1.0286,  1.4379,  1.3279,  0.8112,  1.0858,\n",
      "         0.2515,  1.4586,  1.0680,  1.1439,  0.7283,  0.8081,  0.8144,  1.2165,\n",
      "         1.0143,  0.8507,  0.7017,  1.1180,  0.3363,  1.0810,  0.5907,  1.2289,\n",
      "         0.7880,  1.0597,  1.3484, -0.3572, -0.2874,  0.5995, -0.5068,  1.1236,\n",
      "         1.0205,  1.1258,  1.2615,  0.6402, -0.0742,  1.4522,  1.2662,  0.3906],\n",
      "       device='cuda:0')), ('_final_conv.weight', tensor([[[[ 1.0073,  0.0728,  0.2654,  ...,  0.2754,  0.2594, -0.4779]],\n",
      "\n",
      "         [[ 1.1546, -0.0376,  0.3495,  ...,  0.5901,  0.4206, -0.4803]],\n",
      "\n",
      "         [[ 0.9357,  0.1147,  0.2512,  ...,  0.4635,  0.4616, -0.5045]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 2.0620,  1.6083, -0.6575,  ...,  0.4182,  0.6070, -0.5047]],\n",
      "\n",
      "         [[ 2.6384,  2.5093, -1.1375,  ...,  0.4076,  0.7466, -0.4836]],\n",
      "\n",
      "         [[ 4.5800,  2.7505, -0.8211,  ...,  0.5173,  0.8259,  2.4796]]],\n",
      "\n",
      "\n",
      "        [[[ 0.7975, -0.1232,  0.3188,  ...,  0.2658,  0.3081, -0.4613]],\n",
      "\n",
      "         [[ 1.0612,  0.1709,  0.3675,  ...,  0.5586,  0.3886, -0.5696]],\n",
      "\n",
      "         [[ 0.9073,  0.3344,  0.3498,  ...,  0.4124,  0.4507, -0.4557]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.8232,  1.5379, -0.3717,  ...,  0.3589,  0.5713, -0.5520]],\n",
      "\n",
      "         [[ 2.1677,  2.2583, -0.7822,  ...,  0.3843,  0.6717, -0.5525]],\n",
      "\n",
      "         [[ 3.8519,  2.3095, -0.5379,  ...,  0.5119,  0.6438,  1.9785]]],\n",
      "\n",
      "\n",
      "        [[[ 0.5053, -0.2000,  0.4039,  ...,  0.3719,  0.2937, -0.5170]],\n",
      "\n",
      "         [[ 0.7660,  0.0449,  0.6248,  ...,  0.5970,  0.3926, -0.5808]],\n",
      "\n",
      "         [[ 0.8928,  0.4724,  0.3012,  ...,  0.4167,  0.4238, -0.5712]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.7230,  1.6247, -0.2732,  ...,  0.4026,  0.5962, -0.5628]],\n",
      "\n",
      "         [[ 2.0030,  2.1250, -0.6040,  ...,  0.3594,  0.6040, -0.4700]],\n",
      "\n",
      "         [[ 3.2470,  2.0534, -0.3408,  ...,  0.4919,  0.6108,  1.5860]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 0.4767,  0.6522,  0.5509,  ...,  0.2840,  0.0769, -0.3087]],\n",
      "\n",
      "         [[ 0.3704,  0.3084,  1.0082,  ...,  0.4787,  0.2022, -0.4479]],\n",
      "\n",
      "         [[ 0.3275,  0.4826,  0.8969,  ...,  0.5094,  0.2976, -0.4492]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.1389,  1.3948,  0.8023,  ...,  0.2458,  0.3333, -0.6339]],\n",
      "\n",
      "         [[ 1.2633,  1.5494,  0.7362,  ...,  0.2034,  0.3921, -0.6184]],\n",
      "\n",
      "         [[ 1.4134,  1.0604,  0.7746,  ...,  0.0718,  0.3761,  0.3056]]],\n",
      "\n",
      "\n",
      "        [[[ 0.7007,  0.8919,  0.6103,  ...,  0.3452,  0.1151, -0.2301]],\n",
      "\n",
      "         [[ 0.3520,  0.3877,  1.1085,  ...,  0.5023,  0.2605, -0.4996]],\n",
      "\n",
      "         [[ 0.3159,  0.5195,  1.0100,  ...,  0.4478,  0.2036, -0.4623]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.9656,  1.2825,  1.0610,  ...,  0.3517,  0.2921, -0.6212]],\n",
      "\n",
      "         [[ 1.3682,  1.5815,  0.8783,  ...,  0.2369,  0.3575, -0.5796]],\n",
      "\n",
      "         [[ 1.4736,  1.1305,  0.8715,  ...,  0.0598,  0.4205,  0.2755]]],\n",
      "\n",
      "\n",
      "        [[[ 0.7687,  1.1161,  0.6325,  ...,  0.3290,  0.1120, -0.2019]],\n",
      "\n",
      "         [[ 0.4945,  0.5664,  1.0189,  ...,  0.3566,  0.2002, -0.3913]],\n",
      "\n",
      "         [[ 0.3214,  0.5366,  1.0924,  ...,  0.5289,  0.2116, -0.4167]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.8209,  0.9882,  1.2170,  ...,  0.4113,  0.3042, -0.6328]],\n",
      "\n",
      "         [[ 1.2669,  1.3911,  1.0897,  ...,  0.2229,  0.3691, -0.6365]],\n",
      "\n",
      "         [[ 1.6519,  1.2468,  0.9798,  ...,  0.0105,  0.3706,  0.2052]]]],\n",
      "       device='cuda:0')), ('_final_conv.bias', tensor([1.9253, 1.4408, 1.0507, 1.6709, 1.3665, 1.0312, 1.6930, 1.3336, 1.5856,\n",
      "        1.2820, 1.9035, 1.4442], device='cuda:0'))])\n"
     ]
    }
   ],
   "source": [
    "print(Teacher.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "434f6e0e-643f-4cc0-a410-9b5855466d39",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state \t {0: {'step': tensor(19527.), 'exp_avg': tensor([ 1.0126e-03, -5.4395e-05,  4.7144e-04,  9.2481e-04, -1.1797e-03,\n",
      "         2.1716e-03, -1.8800e-04, -1.4909e-03, -5.6754e-04,  9.3604e-04,\n",
      "         2.5555e-03,  8.1660e-04,  2.6403e-03, -4.5370e-04, -6.5035e-05,\n",
      "         5.9047e-04,  2.2746e-03,  2.1456e-03,  2.3619e-03,  3.2418e-03,\n",
      "         1.4282e-04, -1.0995e-03, -9.4018e-04, -6.8637e-04,  1.2438e-03,\n",
      "         9.0418e-04,  1.5194e-03, -8.3317e-04,  1.2811e-03,  6.2705e-04,\n",
      "         9.1821e-04, -1.0021e-03,  2.9116e-03,  7.3326e-04,  7.7511e-04,\n",
      "         1.6662e-03,  2.1342e-03,  1.1616e-03,  3.9137e-04,  4.0767e-04,\n",
      "         1.5526e-03,  5.3927e-03,  3.1605e-03,  2.9166e-03, -1.5444e-03,\n",
      "         2.4302e-03, -1.3441e-03, -4.5856e-04,  1.1417e-03, -8.8680e-04,\n",
      "        -5.8099e-05,  3.3293e-03,  3.7254e-04,  1.4927e-03,  6.3597e-04,\n",
      "         1.3350e-03,  3.1470e-03,  4.1953e-03,  5.3009e-03,  2.3086e-03,\n",
      "        -1.9088e-03,  2.9275e-03, -6.1441e-04,  1.9771e-03,  7.5513e-04,\n",
      "         1.2770e-03,  8.7307e-04,  2.5871e-03, -1.2065e-03,  3.9824e-03,\n",
      "        -8.5351e-04,  2.0625e-03,  2.7532e-03,  1.6367e-03, -9.5220e-04,\n",
      "         1.4162e-03,  1.9327e-03,  1.3927e-03,  3.1238e-03,  2.8649e-03,\n",
      "         2.0923e-03,  2.1710e-03,  3.1987e-03, -1.5724e-03,  1.2602e-03,\n",
      "         2.0267e-03, -9.5108e-04, -1.1833e-03,  1.7521e-03, -1.6459e-04,\n",
      "         1.0936e-03,  9.6771e-04,  2.4395e-03, -3.5933e-04,  2.0595e-03,\n",
      "         3.7632e-03,  2.8733e-03,  4.3734e-03, -8.4658e-04,  1.3871e-03,\n",
      "         1.4724e-03,  3.7075e-03,  2.8705e-04,  1.3930e-03, -4.5181e-04,\n",
      "        -1.1998e-04,  2.1673e-03, -1.4886e-03,  4.4956e-03,  2.3521e-03,\n",
      "         2.9099e-03,  5.0817e-03, -6.9631e-04, -6.9450e-04,  4.9735e-04,\n",
      "        -6.6464e-04, -3.4911e-04,  6.2620e-04,  1.4018e-03,  3.9411e-03,\n",
      "         3.7431e-03,  2.2678e-03,  3.0506e-03,  1.7014e-03,  2.8475e-03,\n",
      "         1.9190e-03, -6.6823e-04,  2.5855e-03, -1.6715e-03, -1.6587e-03,\n",
      "        -1.7978e-03, -1.5637e-03, -1.6596e-03, -5.8592e-04, -2.8698e-04,\n",
      "         1.6939e-03,  9.4491e-04,  2.0753e-03, -9.3602e-05, -2.7781e-04,\n",
      "         5.3929e-04,  1.0326e-03,  3.8257e-03,  1.2438e-04, -3.6669e-04,\n",
      "        -1.3495e-03,  3.1880e-03,  4.6506e-06,  6.1575e-04,  7.4299e-04,\n",
      "         8.1286e-04,  4.6348e-04, -1.2205e-03,  3.1745e-03,  2.3547e-03,\n",
      "        -1.2624e-03, -1.1283e-03, -9.2716e-04, -3.5631e-04,  2.0097e-03,\n",
      "         1.3676e-03,  7.6816e-06,  1.8954e-03,  1.7911e-03,  2.8301e-03,\n",
      "        -7.6198e-04, -2.3851e-05,  3.2483e-04,  1.4532e-03,  1.9918e-03,\n",
      "         1.9925e-03, -1.9923e-04, -3.6694e-04,  1.9682e-03, -7.3929e-06,\n",
      "         1.9470e-03,  3.0954e-03,  1.9921e-03,  2.1486e-03,  1.9918e-03,\n",
      "        -3.6634e-04, -3.6666e-04, -9.4464e-04, -3.6644e-04, -7.0115e-04,\n",
      "        -5.7889e-04, -1.2433e-03,  1.0892e-04,  2.5882e-04, -3.6662e-04,\n",
      "        -1.3014e-03, -3.6695e-04, -3.6679e-04, -3.6695e-04, -3.6635e-04,\n",
      "        -3.6635e-04, -1.2949e-03, -3.6652e-04, -3.6647e-04, -3.6666e-04,\n",
      "        -3.6663e-04,  1.0205e-03,  1.5576e-03,  4.6476e-03,  2.0578e-03,\n",
      "        -1.2996e-03, -5.5772e-04,  2.8169e-03,  4.9467e-03,  3.2707e-03,\n",
      "         1.9895e-04,  2.3460e-04, -6.1498e-04,  3.8236e-03, -2.4228e-04,\n",
      "         2.5710e-04,  1.5764e-05, -1.2515e-04, -4.8738e-05, -1.7782e-05,\n",
      "         9.9691e-04,  1.3183e-04, -2.5144e-04,  2.7929e-04, -1.4716e-03,\n",
      "        -1.2635e-03, -1.5962e-03, -1.9246e-03, -9.4164e-04, -1.3152e-03,\n",
      "        -5.1860e-04,  6.1648e-05, -9.0533e-04,  7.6147e-05,  3.2984e-03,\n",
      "         1.2479e-03,  2.0172e-03, -2.5741e-04,  1.0565e-03,  2.3247e-03,\n",
      "         1.0860e-03,  8.8325e-04, -2.5507e-04, -1.7739e-04, -5.9823e-04,\n",
      "         6.8769e-04, -1.9460e-04,  2.2656e-05, -8.6465e-04, -1.2951e-03,\n",
      "         2.2965e-03,  8.0566e-04,  2.9262e-03,  3.3283e-03,  7.5751e-04,\n",
      "         4.7669e-04,  3.9615e-04, -1.4636e-03, -1.0641e-03,  3.8535e-04,\n",
      "         1.4718e-03,  1.9921e-03,  1.6467e-03,  1.7946e-03, -4.1495e-04,\n",
      "         9.5151e-04, -1.8653e-03,  1.2099e-03,  1.0045e-03,  3.5328e-06,\n",
      "         2.9025e-03,  2.9332e-03,  4.1021e-03, -1.7243e-03, -1.2796e-03,\n",
      "        -1.4226e-03, -1.5865e-03, -7.8836e-04, -4.5206e-04, -1.5697e-04,\n",
      "        -9.4218e-04,  6.4160e-04, -4.2196e-04,  1.1976e-03,  6.8585e-04,\n",
      "         6.3019e-04,  2.2126e-04,  1.8607e-03,  1.6573e-03, -8.6734e-04,\n",
      "         3.4359e-04,  2.2315e-03,  7.7233e-04,  1.1980e-03,  1.2115e-03,\n",
      "         1.8790e-03, -6.1961e-04,  4.2581e-03,  2.7690e-03,  2.9402e-03,\n",
      "         7.1170e-04,  6.5254e-04,  3.3775e-04,  8.7786e-04,  2.4161e-03,\n",
      "         3.0072e-04, -2.1137e-04], device='cuda:0'), 'exp_avg_sq': tensor([1.0326e-05, 8.6804e-06, 2.4953e-05, 1.2872e-05, 8.5823e-06, 3.6928e-05,\n",
      "        2.3604e-05, 1.8095e-05, 3.5172e-06, 9.7509e-06, 4.0818e-05, 6.7514e-06,\n",
      "        8.6776e-05, 1.3443e-05, 1.4125e-05, 3.0195e-05, 4.3750e-05, 5.3911e-05,\n",
      "        3.8133e-05, 6.3490e-05, 1.2876e-06, 9.0231e-06, 1.2104e-05, 8.8931e-06,\n",
      "        1.3756e-05, 1.4503e-05, 2.0084e-05, 1.5298e-05, 1.4354e-05, 7.4256e-06,\n",
      "        8.6663e-06, 1.6521e-05, 5.3953e-05, 9.7992e-06, 1.6256e-05, 2.4823e-05,\n",
      "        3.2543e-05, 1.8817e-05, 6.1100e-06, 1.7259e-06, 2.1521e-05, 2.0750e-04,\n",
      "        4.9249e-05, 4.9305e-05, 2.2042e-05, 5.5599e-05, 1.7545e-05, 2.4247e-06,\n",
      "        3.5791e-05, 9.1970e-06, 6.6858e-06, 5.6023e-05, 1.1721e-05, 3.4575e-05,\n",
      "        1.0933e-05, 1.6435e-05, 6.8205e-05, 1.2776e-04, 1.4702e-04, 3.4936e-05,\n",
      "        3.0238e-05, 4.4769e-05, 4.9476e-06, 5.0482e-05, 6.5608e-06, 3.0238e-05,\n",
      "        1.1629e-05, 5.5259e-05, 1.6103e-05, 1.3697e-04, 1.5104e-05, 3.4333e-05,\n",
      "        5.0234e-05, 2.2367e-05, 7.8570e-06, 5.5097e-05, 4.5899e-05, 3.5066e-05,\n",
      "        6.7296e-05, 9.5420e-05, 2.3475e-05, 6.3049e-05, 1.3411e-04, 1.7963e-05,\n",
      "        2.2056e-05, 3.7101e-05, 7.6399e-06, 9.5278e-06, 2.5399e-05, 3.7550e-06,\n",
      "        2.6102e-05, 4.6165e-05, 6.9243e-05, 3.3770e-06, 5.2107e-05, 1.0196e-04,\n",
      "        3.6776e-05, 1.9842e-04, 1.3480e-05, 1.8667e-05, 3.8991e-05, 1.4636e-04,\n",
      "        1.3598e-05, 2.1197e-05, 6.6716e-06, 9.4368e-06, 6.2226e-05, 2.6121e-05,\n",
      "        1.9020e-04, 1.0544e-04, 8.9325e-05, 1.8077e-04, 1.2060e-05, 5.2421e-06,\n",
      "        2.2848e-05, 4.9534e-06, 1.2216e-05, 8.0400e-06, 1.3207e-05, 1.2276e-04,\n",
      "        1.1426e-04, 5.8909e-05, 6.2372e-05, 4.8681e-05, 5.4702e-05, 2.7561e-05,\n",
      "        6.3517e-06, 3.9826e-05, 3.2689e-05, 2.0635e-05, 2.5027e-05, 2.5510e-05,\n",
      "        2.8575e-05, 6.6277e-06, 2.4596e-06, 2.0876e-05, 2.1098e-05, 4.3681e-05,\n",
      "        7.0687e-07, 7.0877e-06, 6.7039e-06, 1.2730e-05, 1.2115e-04, 3.4304e-06,\n",
      "        6.3170e-06, 1.6760e-05, 1.0632e-04, 1.7821e-06, 8.9401e-06, 8.2354e-06,\n",
      "        7.5167e-06, 3.0388e-05, 7.9481e-06, 7.2095e-05, 5.1473e-05, 1.0647e-05,\n",
      "        1.1642e-05, 5.0403e-06, 1.9786e-06, 3.0363e-05, 1.3507e-05, 5.5108e-06,\n",
      "        2.3428e-05, 2.6515e-05, 5.6447e-05, 4.4514e-06, 3.8954e-06, 6.4872e-06,\n",
      "        1.5091e-05, 2.9426e-05, 2.9426e-05, 4.5134e-06, 6.3170e-06, 4.3571e-05,\n",
      "        9.0877e-06, 2.4747e-05, 6.7574e-05, 2.9426e-05, 3.1108e-05, 2.9426e-05,\n",
      "        6.3170e-06, 6.3170e-06, 4.6801e-06, 6.3170e-06, 6.4716e-06, 8.6606e-06,\n",
      "        1.1095e-05, 2.3301e-06, 3.1763e-06, 6.3170e-06, 1.5221e-05, 6.3170e-06,\n",
      "        6.3170e-06, 6.3170e-06, 6.3170e-06, 6.3170e-06, 2.0100e-05, 6.3170e-06,\n",
      "        6.3170e-06, 6.3170e-06, 6.3170e-06, 3.0223e-05, 2.4343e-05, 1.8739e-04,\n",
      "        7.1956e-05, 2.2336e-05, 5.8703e-06, 5.0723e-05, 1.3985e-04, 1.2248e-04,\n",
      "        5.8386e-06, 4.5514e-06, 1.0708e-05, 9.3064e-05, 4.9381e-06, 1.6264e-05,\n",
      "        1.5046e-05, 7.6361e-06, 1.3817e-05, 7.7822e-06, 3.4895e-05, 8.2389e-06,\n",
      "        7.6082e-06, 6.2757e-06, 1.2311e-05, 1.7715e-05, 2.7440e-05, 2.5153e-05,\n",
      "        1.1969e-05, 1.0261e-05, 1.8648e-06, 6.1438e-06, 1.3805e-05, 2.9612e-06,\n",
      "        6.8300e-05, 1.3554e-05, 2.3777e-05, 1.8977e-06, 1.7559e-05, 4.6690e-05,\n",
      "        1.1128e-05, 1.8759e-05, 3.4315e-06, 4.3442e-06, 6.8862e-06, 2.6892e-05,\n",
      "        2.8383e-06, 8.4276e-06, 4.8735e-06, 1.3513e-05, 3.7087e-05, 8.4721e-06,\n",
      "        6.0796e-05, 7.5007e-05, 7.3882e-06, 1.2162e-05, 1.4842e-05, 1.9739e-05,\n",
      "        8.3190e-06, 4.1959e-06, 2.7485e-05, 2.9426e-05, 3.0038e-05, 1.8348e-05,\n",
      "        3.2259e-06, 1.4006e-05, 4.2405e-05, 1.5322e-05, 8.6212e-06, 4.3788e-06,\n",
      "        4.4997e-05, 6.2136e-05, 1.4607e-04, 2.1754e-05, 1.6498e-05, 2.1502e-05,\n",
      "        2.0956e-05, 1.1700e-05, 6.2315e-06, 3.8543e-06, 8.9325e-06, 9.7742e-06,\n",
      "        5.8857e-06, 1.2563e-05, 2.0585e-05, 2.7417e-05, 1.1192e-05, 2.4289e-05,\n",
      "        2.1899e-05, 1.2952e-05, 4.3779e-06, 6.8017e-05, 9.3194e-06, 2.1567e-05,\n",
      "        2.3513e-05, 2.8314e-05, 7.3290e-06, 1.1849e-04, 1.1349e-04, 8.0211e-05,\n",
      "        1.3330e-05, 9.6055e-06, 9.0239e-06, 1.8440e-05, 5.8669e-05, 8.4757e-06,\n",
      "        1.0451e-06], device='cuda:0')}, 1: {'step': tensor(19527.), 'exp_avg': tensor([[ 3.7423e-04,  1.4004e-03, -2.6063e-04,  5.7590e-04,  1.3559e-03,\n",
      "         -2.5526e-03,  1.8643e-03,  1.7849e-03,  1.1523e-03,  3.7579e-04,\n",
      "          1.1739e-04, -3.9961e-04, -1.8914e-03,  1.3078e-03,  9.5097e-04,\n",
      "         -1.1587e-04, -1.8950e-03, -1.0520e-03, -4.4340e-04, -2.1325e-03,\n",
      "          6.2060e-05,  1.6808e-03,  2.6714e-03,  2.2256e-03,  2.0495e-04,\n",
      "         -7.1090e-04, -4.6950e-04,  3.2715e-03, -4.2504e-05,  1.4724e-04,\n",
      "         -9.6981e-04,  2.5862e-03, -1.5776e-03,  4.2570e-04, -6.6878e-04,\n",
      "         -1.3305e-03, -3.2810e-04,  4.7070e-04,  7.4000e-04,  2.3808e-04,\n",
      "         -1.5443e-03, -5.4363e-03, -2.5931e-03, -2.4751e-03,  1.8974e-03,\n",
      "         -1.3985e-03,  2.1527e-03,  1.1902e-03, -9.4369e-04,  1.9012e-03,\n",
      "          1.1544e-03, -4.3124e-03, -2.3304e-04, -1.2309e-03,  9.1042e-04,\n",
      "         -1.2978e-03, -2.3463e-03, -5.2226e-03, -1.0620e-02, -7.6211e-04,\n",
      "          2.1476e-03, -2.3259e-03,  1.4803e-03, -1.3071e-03,  9.6299e-04,\n",
      "         -1.2247e-03,  5.4961e-04, -1.9031e-03,  1.2795e-03, -3.8098e-03,\n",
      "          1.2183e-03, -8.9849e-04, -1.5480e-03, -1.4029e-03,  1.7634e-03,\n",
      "         -6.2182e-04,  2.6258e-04, -1.4733e-03, -3.9098e-03, -3.2769e-03,\n",
      "         -1.7380e-03, -2.0221e-03, -3.6374e-03,  1.8128e-03,  1.0167e-04,\n",
      "         -1.9764e-03,  1.2794e-03,  1.9754e-03, -4.1491e-04,  1.0386e-03,\n",
      "         -7.2550e-04, -2.3689e-03, -2.7339e-03,  7.5325e-04, -1.8397e-03,\n",
      "         -1.1998e-03, -7.2595e-04, -2.1209e-03,  1.7325e-03, -1.2199e-03,\n",
      "         -1.0872e-04, -2.1203e-03, -2.2853e-05, -2.5605e-04,  1.3685e-03,\n",
      "          1.0459e-03, -6.8074e-04,  2.1733e-03, -3.6262e-03, -1.3272e-03,\n",
      "         -7.3967e-04, -5.0621e-03,  1.0262e-03,  6.0681e-04,  4.3515e-04,\n",
      "          1.6991e-03,  1.2973e-03,  1.1771e-04, -8.1119e-04, -3.4826e-03,\n",
      "         -3.2810e-03,  1.3546e-04, -1.7958e-03, -1.6421e-04, -2.9747e-03,\n",
      "         -1.5800e-03,  1.8408e-04, -1.7586e-03,  2.1183e-03,  1.8469e-03,\n",
      "          2.3892e-03,  1.9024e-03,  2.1821e-03,  1.6387e-03,  6.3426e-04,\n",
      "         -4.6732e-04,  1.9505e-04, -1.4425e-03,  4.8952e-04,  1.0922e-03,\n",
      "         -3.6326e-04, -3.0070e-04, -3.9548e-03,  5.7169e-04,  4.4639e-04,\n",
      "          1.8295e-03, -3.0358e-03,  6.7796e-04, -8.5395e-04,  3.0445e-04,\n",
      "         -8.6539e-04,  5.7437e-04,  2.5122e-03, -3.4932e-03, -1.5067e-03,\n",
      "          1.5310e-03,  1.8947e-03,  9.2671e-04,  5.3103e-04, -8.3221e-04,\n",
      "         -5.7259e-04, -5.4710e-04, -1.1148e-03, -2.0752e-03, -3.8212e-03,\n",
      "          1.4893e-03,  6.1009e-04, -9.0094e-05, -5.0416e-04, -2.1820e-03,\n",
      "         -2.1820e-03,  3.6861e-04,  4.4643e-04, -7.7848e-04,  1.3092e-03,\n",
      "         -1.0344e-03, -3.2150e-03, -2.1820e-03, -2.5624e-03, -2.1819e-03,\n",
      "          4.4637e-04,  4.4643e-04,  9.6734e-04,  4.4637e-04,  1.1155e-03,\n",
      "          9.5794e-04,  2.0006e-03,  4.9147e-04, -4.6028e-04,  4.4644e-04,\n",
      "          2.0769e-03,  4.4638e-04,  4.4638e-04,  4.4644e-04,  4.4642e-04,\n",
      "          4.4644e-04,  1.6441e-03,  4.4638e-04,  4.4638e-04,  4.4639e-04,\n",
      "          4.4639e-04, -3.7535e-04, -3.1664e-03, -2.8838e-03, -1.4907e-03,\n",
      "          2.0354e-03,  8.3411e-04, -2.3706e-03, -2.5329e-03, -3.8219e-03,\n",
      "          1.0078e-03,  7.1825e-04,  1.8241e-03, -2.3840e-03,  1.9756e-03,\n",
      "          2.4978e-03,  2.2665e-03,  2.4549e-03,  1.7033e-03,  1.9904e-03,\n",
      "         -4.5028e-04,  7.0980e-04,  5.2475e-04,  8.4866e-04,  1.7779e-03,\n",
      "          2.3154e-03,  2.7291e-03,  2.1734e-03,  1.2571e-03,  1.8527e-03,\n",
      "          6.8884e-04,  7.7515e-04,  1.9299e-03,  2.6716e-04, -3.4697e-03,\n",
      "         -8.5617e-04, -2.2393e-03,  1.1987e-04,  5.9443e-04, -6.7552e-04,\n",
      "         -5.9636e-04, -8.2398e-05,  6.9645e-04,  9.6371e-04,  8.5464e-04,\n",
      "         -2.6359e-04,  4.7252e-04,  4.9619e-04,  1.0156e-03,  2.5926e-03,\n",
      "         -4.8792e-04, -1.1295e-03, -3.1199e-03, -2.0444e-03, -8.6518e-04,\n",
      "          1.1020e-03, -5.2443e-04,  2.4958e-03,  1.4313e-03,  5.4039e-04,\n",
      "         -1.1684e-03, -2.1820e-03, -3.1491e-04, -7.9288e-04,  3.3607e-04,\n",
      "         -5.5703e-04,  2.6695e-03,  4.4336e-04, -6.2915e-05,  1.3083e-03,\n",
      "         -2.2577e-03, -9.0676e-04, -4.8807e-03,  2.6434e-03,  9.4570e-04,\n",
      "          2.1543e-03,  1.9735e-03,  1.2922e-03,  7.5222e-04,  1.9393e-04,\n",
      "          1.6369e-03,  1.5960e-03,  1.7279e-03,  7.3945e-04, -1.5797e-04,\n",
      "         -6.3658e-04, -6.1219e-05,  1.5306e-04, -7.2669e-04,  1.2366e-03,\n",
      "          1.5160e-03, -2.1387e-03, -1.9923e-04,  1.9031e-04, -8.1589e-04,\n",
      "         -3.8872e-03,  1.8883e-03, -1.7256e-03, -2.9356e-03, -3.0489e-03,\n",
      "         -3.0677e-05, -4.7099e-04,  6.0771e-04, -8.2702e-04, -1.7823e-03,\n",
      "          8.8740e-04,  5.4280e-04]], device='cuda:0'), 'exp_avg_sq': tensor([[2.4279e-04, 4.3725e-05, 7.5103e-05, 3.0866e-04, 1.7656e-04, 3.6338e-04,\n",
      "         2.8955e-04, 2.2113e-04, 7.3605e-05, 1.7791e-04, 8.2619e-04, 1.0129e-04,\n",
      "         1.0011e-03, 5.0512e-05, 1.3666e-04, 1.4277e-04, 5.4939e-04, 5.4768e-04,\n",
      "         8.6706e-04, 9.9552e-04, 4.7822e-05, 1.1739e-04, 7.1060e-05, 4.7103e-05,\n",
      "         3.1903e-04, 1.5012e-04, 5.3761e-04, 2.0298e-04, 2.4764e-04, 1.0347e-04,\n",
      "         1.5779e-04, 1.4894e-04, 9.8396e-04, 2.2401e-04, 6.2273e-05, 3.5203e-04,\n",
      "         7.4973e-04, 3.0423e-04, 1.3710e-04, 2.4172e-05, 3.0874e-04, 2.1515e-03,\n",
      "         7.2292e-04, 9.3428e-04, 2.3473e-04, 7.7134e-04, 1.1551e-04, 8.0593e-05,\n",
      "         1.5957e-04, 8.8217e-05, 4.1848e-05, 6.5835e-04, 7.5297e-05, 4.1110e-04,\n",
      "         1.7803e-04, 3.4432e-04, 1.1317e-03, 1.3621e-03, 2.0167e-03, 6.5746e-04,\n",
      "         4.0134e-04, 7.4419e-04, 5.8172e-05, 4.0098e-04, 1.7741e-04, 3.4025e-04,\n",
      "         3.0545e-04, 7.3512e-04, 1.5283e-04, 1.6032e-03, 1.0762e-04, 6.5089e-04,\n",
      "         7.0073e-04, 2.1245e-04, 1.3675e-04, 3.2658e-04, 9.1970e-04, 1.9260e-04,\n",
      "         1.1429e-03, 6.2898e-04, 4.5635e-04, 3.8852e-04, 1.0516e-03, 2.4821e-04,\n",
      "         3.7612e-04, 7.0166e-04, 9.5748e-05, 1.3657e-04, 5.4541e-04, 2.3633e-05,\n",
      "         1.1587e-04, 5.1094e-04, 6.5403e-04, 6.3538e-05, 6.2210e-04, 1.3268e-03,\n",
      "         6.0639e-04, 3.2729e-03, 6.7715e-05, 1.6965e-04, 9.0540e-04, 2.0507e-03,\n",
      "         7.5533e-05, 4.1252e-04, 4.3669e-05, 8.4790e-05, 6.2961e-04, 1.8994e-04,\n",
      "         2.2894e-03, 5.2870e-04, 1.5429e-03, 2.8600e-03, 3.3575e-05, 4.1135e-05,\n",
      "         7.7270e-05, 9.1298e-05, 1.4842e-04, 3.0654e-04, 3.7227e-04, 1.2717e-03,\n",
      "         1.5710e-03, 1.0228e-03, 8.3273e-04, 6.9521e-04, 7.3120e-04, 3.3659e-04,\n",
      "         1.3699e-04, 6.0017e-04, 2.4597e-04, 2.4855e-04, 3.0885e-04, 2.1405e-04,\n",
      "         2.3679e-04, 5.1465e-05, 3.7464e-05, 4.3560e-04, 5.0980e-04, 6.4696e-04,\n",
      "         2.1492e-05, 4.5936e-05, 1.4797e-04, 2.4902e-04, 1.2984e-03, 4.9687e-05,\n",
      "         3.9147e-05, 1.3695e-04, 1.2094e-03, 2.7547e-05, 9.8608e-05, 1.0190e-04,\n",
      "         9.4540e-05, 1.6194e-04, 8.5095e-05, 9.2493e-04, 6.4556e-04, 1.5675e-04,\n",
      "         9.7477e-05, 1.1559e-04, 3.8874e-05, 6.9197e-04, 3.3804e-04, 3.2568e-05,\n",
      "         4.7497e-04, 2.8061e-04, 8.0985e-04, 1.1093e-04, 7.2876e-05, 5.7243e-05,\n",
      "         2.9781e-04, 3.8996e-04, 3.8996e-04, 3.8326e-05, 3.9147e-05, 7.1262e-04,\n",
      "         8.1586e-05, 3.6618e-04, 8.3942e-04, 3.8996e-04, 4.3186e-04, 3.8996e-04,\n",
      "         3.9147e-05, 3.9147e-05, 9.0505e-05, 3.9147e-05, 4.5988e-05, 7.6593e-05,\n",
      "         1.1027e-04, 8.6160e-05, 3.7409e-05, 3.9147e-05, 1.5907e-04, 3.9147e-05,\n",
      "         3.9147e-05, 3.9147e-05, 3.9147e-05, 3.9147e-05, 1.5096e-04, 3.9147e-05,\n",
      "         3.9147e-05, 3.9147e-05, 3.9147e-05, 1.7103e-04, 3.6779e-04, 2.8913e-03,\n",
      "         4.2124e-04, 1.5661e-04, 4.6594e-05, 8.1192e-04, 2.5495e-03, 8.8826e-04,\n",
      "         1.7517e-04, 8.3464e-05, 4.1539e-05, 1.5375e-03, 1.7682e-04, 6.3879e-04,\n",
      "         5.8845e-04, 3.0769e-04, 5.2644e-04, 3.3043e-04, 2.6683e-04, 1.1665e-04,\n",
      "         3.6212e-05, 1.4008e-04, 2.3090e-04, 1.1698e-04, 1.4932e-04, 3.2158e-04,\n",
      "         8.6363e-05, 1.8384e-04, 6.1536e-05, 6.3019e-05, 8.2139e-05, 4.1936e-05,\n",
      "         1.3071e-03, 1.9898e-04, 4.3444e-04, 4.8025e-05, 2.4174e-04, 1.0523e-03,\n",
      "         1.3176e-04, 2.1745e-04, 5.3040e-05, 7.6182e-05, 8.0858e-05, 1.0425e-04,\n",
      "         6.2062e-05, 5.1918e-05, 1.0189e-04, 9.6959e-05, 5.1360e-04, 1.0184e-04,\n",
      "         6.6995e-04, 1.1634e-03, 1.0928e-04, 2.2778e-04, 4.6283e-05, 1.6420e-04,\n",
      "         1.1996e-04, 9.4833e-05, 2.8687e-04, 3.8996e-04, 6.8639e-04, 3.3744e-04,\n",
      "         5.6628e-05, 1.5275e-04, 2.9020e-04, 4.1204e-04, 2.1917e-04, 5.7655e-05,\n",
      "         8.2612e-04, 6.3662e-04, 1.4275e-03, 2.1114e-04, 1.4972e-04, 1.6561e-04,\n",
      "         2.7680e-04, 1.3971e-04, 6.1722e-05, 7.3642e-05, 1.0583e-04, 2.1544e-04,\n",
      "         4.7100e-05, 3.4217e-04, 1.1327e-04, 1.4890e-04, 6.8840e-05, 6.4871e-04,\n",
      "         4.4150e-04, 1.1747e-04, 8.5120e-05, 4.5483e-04, 7.3148e-05, 3.2010e-04,\n",
      "         1.6426e-04, 3.5139e-04, 1.1868e-04, 2.3943e-03, 5.1732e-04, 9.5688e-04,\n",
      "         1.3839e-04, 7.5808e-05, 1.2407e-04, 4.7218e-04, 8.3474e-04, 1.8110e-04,\n",
      "         2.7477e-05]], device='cuda:0')}, 2: {'step': tensor(19527.), 'exp_avg': tensor([0.0765], device='cuda:0'), 'exp_avg_sq': tensor([0.3169], device='cuda:0')}, 3: {'step': tensor(19527.), 'exp_avg': tensor([[[-6.3728e-05, -5.4126e-05, -1.2584e-05, -8.6512e-05, -2.5352e-05,\n",
      "          -1.0815e-04, -1.2478e-04, -1.1810e-04,  2.4295e-05,  5.3704e-06,\n",
      "          -1.3444e-04,  4.8184e-05],\n",
      "         [ 8.4532e-06,  1.6163e-05, -2.4164e-05,  3.4872e-05,  1.0039e-05,\n",
      "           5.0103e-05,  5.7962e-05,  3.5889e-05, -1.1695e-05, -3.9781e-06,\n",
      "           1.0131e-04, -1.5562e-05],\n",
      "         [-3.5044e-05, -3.2078e-05,  2.7957e-05, -3.2012e-05,  4.6185e-06,\n",
      "          -8.9676e-05, -6.9680e-05, -8.0816e-05,  2.3462e-05,  4.6207e-05,\n",
      "          -6.1499e-06,  4.2470e-05],\n",
      "         [-5.6878e-05, -4.6397e-05, -8.4222e-07, -2.3002e-05,  3.1109e-06,\n",
      "          -3.0652e-05, -2.9767e-05, -3.0184e-05, -2.8005e-06, -4.8341e-06,\n",
      "          -4.8894e-05,  8.7031e-06],\n",
      "         [-4.1121e-05, -2.6519e-05,  3.1499e-05,  1.5358e-05,  1.0233e-05,\n",
      "           9.4607e-06,  7.7253e-06, -1.1999e-05,  5.5455e-06,  9.5368e-06,\n",
      "           7.2730e-07,  2.7055e-05],\n",
      "         [ 5.8340e-05,  6.1256e-05, -4.1763e-05,  6.7096e-05,  2.7883e-06,\n",
      "           1.4663e-04,  1.0461e-04,  9.8456e-05, -1.5579e-05, -3.3528e-05,\n",
      "           1.3266e-04, -2.3591e-05],\n",
      "         [ 3.7518e-05,  3.0777e-05,  5.6731e-06,  1.4232e-05, -8.9902e-06,\n",
      "           5.1316e-05, -1.4776e-05, -1.2777e-05,  1.3548e-05,  6.7466e-06,\n",
      "           1.9208e-05,  5.1890e-05],\n",
      "         [-2.8226e-06, -1.8345e-06,  3.8131e-05, -5.0979e-05, -2.1973e-05,\n",
      "          -3.1121e-05, -1.0923e-04, -1.0335e-04,  3.1314e-05,  3.2655e-05,\n",
      "          -4.8813e-05,  7.7380e-05],\n",
      "         [-4.7203e-05, -5.0360e-05,  9.6733e-05, -7.3142e-05, -6.7872e-06,\n",
      "          -1.4374e-04, -1.9641e-04, -2.2087e-04,  6.3215e-05,  8.0290e-05,\n",
      "          -1.1482e-04,  1.4912e-04],\n",
      "         [-1.8190e-05, -2.8378e-05,  8.7910e-05, -2.1362e-05,  3.9747e-06,\n",
      "          -4.8004e-05, -8.3995e-05, -9.6515e-05,  4.8602e-05,  7.9560e-05,\n",
      "          -2.1040e-05,  1.2550e-04],\n",
      "         [-8.9608e-05, -1.0445e-04,  1.8818e-04, -9.1467e-05, -6.5055e-06,\n",
      "          -1.5034e-04, -1.8341e-04, -2.1225e-04,  6.6096e-05,  1.0066e-04,\n",
      "          -1.4143e-04,  1.8071e-04],\n",
      "         [-1.6026e-04, -2.0205e-04,  1.5405e-04, -1.5213e-04, -2.7660e-05,\n",
      "          -2.9157e-04, -2.6655e-04, -2.8045e-04,  6.6260e-05,  8.4451e-05,\n",
      "          -3.3887e-04,  1.9144e-04]]], device='cuda:0'), 'exp_avg_sq': tensor([[[3.7401e-07, 2.3835e-07, 3.3203e-07, 2.5749e-07, 1.6610e-08,\n",
      "          8.1501e-07, 4.4315e-07, 6.0579e-07, 6.0582e-08, 1.7427e-07,\n",
      "          4.7781e-07, 2.3273e-07],\n",
      "         [1.3370e-07, 8.4756e-08, 1.2607e-07, 7.8481e-08, 5.9214e-09,\n",
      "          2.5391e-07, 1.3666e-07, 1.8600e-07, 2.0490e-08, 6.5092e-08,\n",
      "          1.3267e-07, 9.1574e-08],\n",
      "         [1.2674e-07, 7.6810e-08, 1.2518e-07, 9.1310e-08, 8.3366e-09,\n",
      "          2.9588e-07, 1.5467e-07, 2.1267e-07, 2.2165e-08, 6.4849e-08,\n",
      "          1.7859e-07, 1.0609e-07],\n",
      "         [6.5412e-08, 4.0290e-08, 4.9233e-08, 4.3224e-08, 3.3917e-09,\n",
      "          1.2057e-07, 6.4946e-08, 9.7037e-08, 1.0092e-08, 3.0820e-08,\n",
      "          7.1889e-08, 3.9411e-08],\n",
      "         [7.7362e-08, 4.8699e-08, 8.2821e-08, 6.4099e-08, 5.6703e-09,\n",
      "          1.8591e-07, 1.0351e-07, 1.3564e-07, 1.6351e-08, 4.7986e-08,\n",
      "          1.1184e-07, 7.1553e-08],\n",
      "         [1.4894e-07, 9.7907e-08, 1.4371e-07, 7.7773e-08, 7.1064e-09,\n",
      "          2.8301e-07, 1.5690e-07, 2.5144e-07, 2.6454e-08, 8.9347e-08,\n",
      "          1.6275e-07, 1.0272e-07],\n",
      "         [2.0715e-07, 1.3763e-07, 1.2394e-07, 1.8295e-07, 1.3835e-08,\n",
      "          5.5175e-07, 3.0608e-07, 3.9441e-07, 2.3251e-08, 7.4480e-08,\n",
      "          3.0192e-07, 8.4052e-08],\n",
      "         [3.7345e-07, 2.3258e-07, 2.2418e-07, 3.5906e-07, 2.6634e-08,\n",
      "          9.6248e-07, 5.0435e-07, 6.1785e-07, 3.8408e-08, 1.1144e-07,\n",
      "          5.4195e-07, 1.3646e-07],\n",
      "         [4.4322e-07, 2.7736e-07, 3.4990e-07, 4.0071e-07, 3.1140e-08,\n",
      "          1.0577e-06, 5.5477e-07, 6.9313e-07, 6.1413e-08, 1.7619e-07,\n",
      "          5.7990e-07, 2.1258e-07],\n",
      "         [1.1712e-07, 6.7882e-08, 1.7469e-07, 1.0508e-07, 1.7958e-08,\n",
      "          2.8212e-07, 1.3833e-07, 1.4882e-07, 2.5499e-08, 8.1071e-08,\n",
      "          1.3801e-07, 9.7644e-08],\n",
      "         [4.0526e-07, 2.3979e-07, 3.7589e-07, 4.4010e-07, 4.0479e-08,\n",
      "          9.9494e-07, 5.0108e-07, 5.3449e-07, 5.3618e-08, 1.5091e-07,\n",
      "          5.4280e-07, 2.1333e-07],\n",
      "         [1.1018e-06, 6.9380e-07, 1.1892e-06, 9.6767e-07, 8.4654e-08,\n",
      "          2.5090e-06, 1.3096e-06, 1.6897e-06, 1.8373e-07, 5.4129e-07,\n",
      "          1.4053e-06, 6.9719e-07]]], device='cuda:0')}, 4: {'step': tensor(19527.), 'exp_avg': tensor([[-3.4436e-05,  4.4081e-05,  3.1831e-04,  1.1948e-04,  3.0757e-05,\n",
      "         -1.3732e-04,  3.8742e-05,  5.3205e-05,  1.5627e-04,  4.8207e-05,\n",
      "          2.1999e-05,  1.2342e-05],\n",
      "        [ 1.9862e-04,  2.0771e-04,  3.0784e-04,  1.3377e-04,  5.5221e-05,\n",
      "         -1.1994e-04,  2.3465e-05,  9.8925e-05,  2.7464e-04,  1.1101e-04,\n",
      "          1.9481e-04,  1.3251e-04],\n",
      "        [-3.9157e-04, -2.4419e-04,  9.7997e-05, -1.1101e-04, -1.4809e-04,\n",
      "         -3.0430e-04, -1.4120e-04, -6.8395e-05,  1.3556e-04,  9.7901e-05,\n",
      "          1.4098e-04,  1.4657e-04],\n",
      "        [ 9.0499e-04,  1.0883e-03,  1.2549e-03,  1.3108e-03,  1.0972e-03,\n",
      "         -6.7972e-05,  5.3718e-05,  2.2135e-04,  5.9819e-04,  6.0091e-04,\n",
      "          7.9380e-04,  2.8034e-04],\n",
      "        [-8.9098e-04, -7.5660e-04, -5.1564e-04, -8.1061e-04, -6.7403e-04,\n",
      "         -6.1190e-04, -5.6891e-04, -5.2053e-04, -5.4710e-05,  1.4542e-04,\n",
      "          2.2043e-04, -3.0514e-04],\n",
      "        [-1.1427e-03, -9.8715e-04, -5.5283e-04, -8.7902e-04, -7.9336e-04,\n",
      "         -6.4211e-04, -5.8442e-04, -4.6852e-04, -1.6597e-04,  6.3573e-05,\n",
      "          1.3712e-04, -7.5751e-04],\n",
      "        [ 2.2482e-03,  2.1021e-03,  2.0938e-03,  2.3761e-03,  2.1820e-03,\n",
      "          1.8754e-03,  2.2101e-03,  2.1925e-03,  1.8923e-03,  1.7247e-03,\n",
      "          1.6663e-03,  7.2682e-04],\n",
      "        [ 3.8021e-03,  3.5657e-03,  3.3589e-03,  3.9020e-03,  3.7349e-03,\n",
      "          3.8825e-03,  4.3781e-03,  4.3832e-03,  3.9395e-03,  3.5957e-03,\n",
      "          3.6874e-03,  2.5263e-03],\n",
      "        [-3.8828e-04, -3.4854e-04, -2.3183e-04, -2.7518e-04, -3.2928e-04,\n",
      "         -1.2493e-05,  3.1163e-04,  2.9315e-04,  1.5866e-04,  1.2842e-04,\n",
      "          8.9155e-05,  1.2118e-03],\n",
      "        [-6.8308e-05, -3.8687e-04, -5.2890e-04, -3.0601e-04, -1.8966e-04,\n",
      "          3.4522e-04,  4.3164e-04,  2.4780e-04, -1.9763e-04, -5.1989e-04,\n",
      "         -7.4890e-04,  6.8233e-04],\n",
      "        [-4.6882e-04, -5.8728e-04, -1.0780e-03, -9.8883e-04, -7.9808e-04,\n",
      "         -3.9313e-04, -1.0661e-03, -1.2814e-03, -1.4768e-03, -1.2275e-03,\n",
      "         -1.1898e-03, -3.4755e-03],\n",
      "        [-3.7687e-03, -3.6965e-03, -4.5245e-03, -4.4701e-03, -4.1664e-03,\n",
      "         -3.8127e-03, -5.0880e-03, -5.1535e-03, -5.2617e-03, -4.7687e-03,\n",
      "         -5.0147e-03, -1.1817e-03]], device='cuda:0'), 'exp_avg_sq': tensor([[7.2401e-06, 3.2023e-06, 1.6911e-06, 1.8492e-06, 1.0621e-06, 8.4030e-07,\n",
      "         1.7566e-06, 1.5910e-06, 1.1593e-06, 7.0835e-07, 1.7640e-06, 2.2309e-06],\n",
      "        [9.7448e-06, 6.0331e-06, 2.9660e-06, 2.7946e-06, 1.2926e-06, 1.0087e-06,\n",
      "         1.8594e-06, 2.0322e-06, 1.3262e-06, 8.7312e-07, 1.8114e-06, 2.5459e-06],\n",
      "        [1.2334e-05, 9.1462e-06, 8.4452e-06, 9.3862e-06, 5.1896e-06, 3.5355e-06,\n",
      "         6.1318e-06, 5.8539e-06, 4.6408e-06, 2.8213e-06, 4.7111e-06, 4.7613e-06],\n",
      "        [1.0791e-04, 1.1410e-04, 1.1606e-04, 1.2898e-04, 9.8403e-05, 8.2186e-05,\n",
      "         8.2637e-05, 7.6094e-05, 6.0089e-05, 4.9398e-05, 5.1912e-05, 4.9131e-05],\n",
      "        [8.3552e-05, 7.7871e-05, 8.1600e-05, 8.7774e-05, 8.1211e-05, 6.8991e-05,\n",
      "         7.7736e-05, 7.1199e-05, 6.2391e-05, 5.3125e-05, 5.6629e-05, 5.4692e-05],\n",
      "        [2.1640e-04, 2.1148e-04, 2.1160e-04, 2.1952e-04, 2.0935e-04, 2.0742e-04,\n",
      "         2.1403e-04, 2.1458e-04, 2.0427e-04, 1.9063e-04, 1.9646e-04, 1.9189e-04],\n",
      "        [2.0544e-04, 1.9644e-04, 1.9472e-04, 1.9943e-04, 1.9813e-04, 1.9647e-04,\n",
      "         2.0955e-04, 2.0980e-04, 2.0114e-04, 1.9191e-04, 2.0154e-04, 1.9844e-04],\n",
      "        [2.2677e-04, 2.1701e-04, 2.1679e-04, 2.2496e-04, 2.1707e-04, 2.1908e-04,\n",
      "         2.3748e-04, 2.4092e-04, 2.2476e-04, 2.1055e-04, 2.1557e-04, 2.0993e-04],\n",
      "        [4.0708e-04, 3.9006e-04, 3.8355e-04, 3.9254e-04, 3.8495e-04, 3.7969e-04,\n",
      "         4.0182e-04, 3.8987e-04, 3.8285e-04, 3.6087e-04, 3.7655e-04, 3.6410e-04],\n",
      "        [3.2439e-04, 3.2183e-04, 3.1622e-04, 3.2842e-04, 3.0774e-04, 3.0669e-04,\n",
      "         3.2862e-04, 3.2572e-04, 3.1077e-04, 2.9272e-04, 3.0288e-04, 2.9337e-04],\n",
      "        [7.2048e-04, 6.7752e-04, 6.8552e-04, 6.9555e-04, 6.8551e-04, 6.7646e-04,\n",
      "         7.0808e-04, 7.0644e-04, 6.7127e-04, 6.4022e-04, 6.6037e-04, 6.5051e-04],\n",
      "        [1.7017e-03, 1.6654e-03, 1.6613e-03, 1.6919e-03, 1.6836e-03, 1.6707e-03,\n",
      "         1.7315e-03, 1.7206e-03, 1.6798e-03, 1.6441e-03, 1.6747e-03, 1.6594e-03]],\n",
      "       device='cuda:0')}, 5: {'step': tensor(19527.), 'exp_avg': tensor([-0.0182, -0.0179, -0.0178, -0.0183, -0.0182, -0.0180, -0.0182, -0.0188,\n",
      "        -0.0186, -0.0188, -0.0192, -0.0186], device='cuda:0'), 'exp_avg_sq': tensor([0.0943, 0.0942, 0.0945, 0.0939, 0.0940, 0.0937, 0.0935, 0.0933, 0.0937,\n",
      "        0.0935, 0.0929, 0.0932], device='cuda:0')}, 6: {'step': tensor(19527.), 'exp_avg': tensor([[0.0431, 0.0442, 0.0428, 0.0446, 0.0437, 0.0454, 0.0455, 0.0456, 0.0441,\n",
      "         0.0439, 0.0478, 0.0453]], device='cuda:0'), 'exp_avg_sq': tensor([[0.6030, 0.6030, 0.6065, 0.6031, 0.6045, 0.6023, 0.6022, 0.6011, 0.6039,\n",
      "         0.6035, 0.5998, 0.6024]], device='cuda:0')}, 7: {'step': tensor(19527.), 'exp_avg': tensor([-0.0825], device='cuda:0'), 'exp_avg_sq': tensor([1.8812], device='cuda:0')}, 8: {'step': tensor(19527.), 'exp_avg': tensor([[[-1.2439e-09, -1.9889e-08, -1.2124e-04,  ..., -5.1108e-09,\n",
      "          -3.9520e-10,  4.5029e-09],\n",
      "         [-1.2766e-08, -2.6512e-07, -6.1342e-06,  ..., -7.3806e-09,\n",
      "          -3.2280e-10,  2.6283e-09],\n",
      "         [-8.7281e-08,  3.6455e-07,  6.0860e-05,  ..., -2.2535e-08,\n",
      "          -1.6053e-09,  7.3128e-09],\n",
      "         ...,\n",
      "         [-1.3289e-08, -8.1231e-08,  1.8055e-05,  ..., -1.9753e-08,\n",
      "          -9.7958e-10,  2.2320e-09],\n",
      "         [ 8.4567e-10, -2.5910e-09,  1.2264e-05,  ...,  1.1038e-10,\n",
      "          -8.0099e-11,  4.9225e-11],\n",
      "         [-3.4815e-09, -1.1798e-08,  1.6001e-05,  ..., -1.8193e-09,\n",
      "          -3.3429e-10,  1.1336e-09]]], device='cuda:0'), 'exp_avg_sq': tensor([[[3.0864e-16, 2.9569e-12, 1.2282e-06,  ..., 1.2458e-13,\n",
      "          1.7783e-16, 3.4733e-16],\n",
      "         [3.2104e-15, 3.0608e-12, 1.0853e-08,  ..., 2.1991e-13,\n",
      "          1.4922e-16, 1.7676e-16],\n",
      "         [7.9491e-14, 7.3401e-12, 1.4174e-07,  ..., 1.7355e-13,\n",
      "          1.5811e-15, 2.4918e-15],\n",
      "         ...,\n",
      "         [2.5522e-15, 1.1766e-12, 3.2856e-08,  ..., 3.8263e-13,\n",
      "          4.9848e-17, 9.3527e-17],\n",
      "         [1.1459e-16, 8.6720e-14, 1.8459e-08,  ..., 9.9277e-15,\n",
      "          1.4015e-18, 4.9418e-18],\n",
      "         [2.0943e-16, 6.5640e-13, 1.6821e-07,  ..., 7.8958e-14,\n",
      "          8.8499e-18, 2.1485e-17]]], device='cuda:0')}, 9: {'step': tensor(19527.), 'exp_avg': tensor([[-4.1629e-07,  2.7465e-07, -1.3599e-07,  ..., -3.9252e-07,\n",
      "         -6.9247e-07, -1.9466e-07],\n",
      "        [ 1.6910e-06,  4.2525e-06,  2.1786e-06,  ...,  1.1759e-06,\n",
      "          1.4844e-06,  2.0304e-06],\n",
      "        [-2.6769e-04, -2.4576e-04, -1.2852e-04,  ..., -1.7673e-04,\n",
      "          1.2364e-04, -2.3228e-04],\n",
      "        ...,\n",
      "        [-1.6436e-08,  1.7742e-07,  7.7349e-08,  ..., -1.7891e-07,\n",
      "         -1.7742e-07,  1.4885e-08],\n",
      "        [-8.8594e-08, -2.3348e-08, -2.2214e-07,  ..., -8.5280e-08,\n",
      "         -1.3723e-07, -6.1897e-08],\n",
      "        [-1.0472e-07, -3.5800e-08, -3.6407e-07,  ..., -9.7464e-08,\n",
      "         -1.1979e-07, -5.9433e-08]], device='cuda:0'), 'exp_avg_sq': tensor([[4.5549e-12, 7.1075e-12, 2.9664e-11,  ..., 2.7135e-12, 5.6087e-12,\n",
      "         4.7321e-12],\n",
      "        [1.2847e-10, 3.1408e-10, 1.6604e-09,  ..., 6.7992e-11, 4.3998e-11,\n",
      "         7.1176e-11],\n",
      "        [2.5980e-05, 1.2761e-05, 1.8778e-06,  ..., 7.3891e-06, 1.8340e-04,\n",
      "         7.6549e-05],\n",
      "        ...,\n",
      "        [3.3770e-11, 3.2086e-11, 1.6443e-10,  ..., 3.2446e-11, 2.0005e-11,\n",
      "         2.2922e-11],\n",
      "        [1.1402e-13, 1.7318e-13, 6.3533e-13,  ..., 7.7740e-14, 1.8339e-13,\n",
      "         9.5534e-14],\n",
      "        [3.2790e-14, 6.0353e-14, 5.5425e-13,  ..., 3.0305e-14, 3.5857e-14,\n",
      "         2.1833e-14]], device='cuda:0')}, 10: {'step': tensor(19527.), 'exp_avg': tensor([[[-5.0276e-05,  3.2060e-04,  3.7426e-03, -1.4000e-03, -5.4021e-05,\n",
      "           3.8958e-05, -5.4759e-04, -2.2250e-04, -5.4322e-04,  4.1709e-04,\n",
      "          -5.1242e-04, -1.0587e-03,  5.4063e-05, -9.1972e-04, -6.2456e-04,\n",
      "           6.3372e-05, -7.7239e-04,  2.2455e-05,  2.0552e-05, -6.0118e-05,\n",
      "          -7.8986e-04,  2.1456e-04,  9.4674e-05,  3.7223e-04,  4.0155e-04,\n",
      "          -4.7391e-05,  1.3707e-04, -7.7426e-04,  2.3184e-04, -1.7914e-03,\n",
      "           2.6411e-05, -8.0002e-04,  7.2802e-05,  2.9272e-03, -1.5538e-03,\n",
      "           2.1513e-05,  6.7580e-04, -6.8848e-05, -1.0902e-03, -8.0827e-05,\n",
      "          -5.0975e-04, -4.4010e-05,  3.3406e-04, -9.8335e-04, -7.5507e-05,\n",
      "          -9.6402e-05, -8.2350e-06, -4.3133e-04, -3.4440e-04, -5.2108e-04,\n",
      "          -4.3968e-04, -7.1447e-04, -1.4822e-05, -8.7272e-05, -5.5583e-04,\n",
      "           2.2214e-04, -7.6581e-05, -6.3009e-04, -1.5324e-03, -7.0209e-04,\n",
      "          -1.7137e-05,  1.4485e-05,  4.1971e-03, -8.5101e-04]],\n",
      "\n",
      "        [[-6.6618e-05,  2.2905e-04, -6.0078e-04,  2.2295e-04, -6.7003e-05,\n",
      "           1.7969e-04,  1.2838e-04,  3.2259e-04,  1.8309e-04,  1.5711e-04,\n",
      "           3.2688e-05, -1.0275e-05, -1.1080e-04, -9.9190e-05, -9.1852e-05,\n",
      "           1.0952e-05, -2.1039e-04,  1.3519e-04,  3.6018e-04,  7.9865e-07,\n",
      "          -1.5741e-04,  1.0677e-05,  1.1144e-05,  9.7344e-06,  1.6350e-04,\n",
      "           2.2753e-06,  7.3071e-06,  2.5970e-04,  2.5791e-04,  8.8514e-05,\n",
      "           3.0638e-04,  3.4961e-04,  1.7489e-04,  5.6144e-04,  7.3498e-05,\n",
      "          -4.6253e-05,  4.8061e-05,  2.3694e-04, -1.3404e-04, -2.9612e-05,\n",
      "           1.8874e-04,  6.7811e-05,  1.6925e-04, -6.2239e-05, -3.9799e-05,\n",
      "           1.5598e-04, -2.0325e-05,  9.6837e-05, -9.4393e-05,  1.1932e-04,\n",
      "           1.0362e-04, -6.1969e-05,  7.7141e-06, -1.0025e-05, -3.5667e-05,\n",
      "          -4.3588e-05, -2.6324e-05,  9.7582e-05, -2.0067e-04, -6.5809e-05,\n",
      "          -1.1955e-06, -5.3253e-05,  7.9084e-04,  7.0237e-05]],\n",
      "\n",
      "        [[ 1.4899e-04, -6.8590e-04, -3.1172e-03,  8.4338e-04,  2.4880e-04,\n",
      "          -3.2451e-04,  2.6967e-04, -3.1178e-04,  2.8541e-04, -5.5994e-04,\n",
      "           4.5515e-04,  7.9420e-04,  9.0875e-05,  1.1084e-03,  6.8786e-04,\n",
      "           4.5038e-05,  1.7337e-03, -2.8826e-04, -3.4187e-04,  1.3976e-05,\n",
      "           9.7175e-04, -2.2608e-04, -1.6112e-05, -3.4270e-04, -5.0938e-04,\n",
      "          -6.9228e-06,  1.7502e-05,  3.6299e-04, -5.2349e-04,  2.5774e-03,\n",
      "          -4.5809e-04, -3.8718e-05, -3.3177e-04, -3.4903e-03,  2.2151e-03,\n",
      "          -2.6623e-05, -7.0622e-04, -2.9248e-04,  1.1714e-03,  4.6499e-05,\n",
      "           2.4897e-04, -1.0414e-04, -4.7301e-04,  8.1683e-04,  2.6043e-04,\n",
      "          -5.3792e-05, -9.6408e-06,  2.7180e-04,  4.6766e-04,  3.0324e-04,\n",
      "           2.9875e-04,  6.3551e-04,  1.2556e-04,  1.9572e-04,  4.4813e-04,\n",
      "          -1.5661e-04,  2.0096e-04,  2.9276e-04,  1.8317e-03,  7.3411e-04,\n",
      "           2.1370e-05,  2.4242e-04, -5.0494e-03,  5.3889e-04]]],\n",
      "       device='cuda:0'), 'exp_avg_sq': tensor([[[2.9052e-05, 1.3358e-04, 5.9870e-05, 2.9274e-05, 8.6075e-06,\n",
      "          6.1049e-05, 1.9018e-05, 1.3372e-04, 2.9869e-05, 7.6216e-05,\n",
      "          3.2500e-06, 2.5730e-05, 1.2977e-04, 5.4606e-05, 4.2529e-05,\n",
      "          1.0325e-05, 1.4288e-04, 4.1938e-05, 2.9135e-04, 2.1136e-06,\n",
      "          9.1342e-05, 1.1229e-05, 1.2268e-05, 6.7032e-06, 8.8999e-05,\n",
      "          1.0954e-06, 1.2261e-05, 6.4903e-05, 1.7403e-04, 4.6806e-05,\n",
      "          1.6553e-04, 1.3923e-04, 6.6234e-05, 1.9345e-03, 3.6298e-05,\n",
      "          1.9947e-05, 4.7985e-05, 7.1393e-05, 9.6707e-05, 8.3143e-06,\n",
      "          4.2505e-05, 9.3279e-06, 7.7813e-05, 2.7586e-05, 1.2155e-05,\n",
      "          4.2906e-05, 1.9328e-06, 6.2228e-06, 6.3012e-05, 1.1899e-05,\n",
      "          5.6605e-06, 2.6026e-05, 1.5681e-07, 2.9801e-06, 1.1566e-05,\n",
      "          2.7076e-05, 4.6968e-06, 9.3924e-06, 1.5786e-04, 2.6982e-05,\n",
      "          8.5083e-07, 2.4086e-05, 4.5838e-03, 7.8197e-06]],\n",
      "\n",
      "        [[3.4000e-07, 1.0388e-06, 7.6155e-07, 1.6995e-07, 3.3988e-07,\n",
      "          4.9467e-07, 1.3570e-07, 1.0876e-06, 2.2763e-07, 5.2979e-07,\n",
      "          7.8586e-09, 1.3018e-07, 1.3834e-06, 3.7532e-07, 3.0959e-07,\n",
      "          2.9221e-08, 5.4166e-06, 3.0650e-07, 2.1899e-06, 9.7629e-09,\n",
      "          7.0154e-07, 4.0970e-08, 2.7510e-08, 3.4430e-08, 6.7817e-07,\n",
      "          3.2717e-09, 4.0029e-08, 4.8047e-07, 1.2638e-06, 3.2963e-07,\n",
      "          1.3041e-06, 1.1355e-06, 5.5495e-07, 1.3693e-05, 2.4162e-07,\n",
      "          2.3369e-07, 3.3387e-07, 5.4975e-07, 7.0964e-07, 9.8177e-08,\n",
      "          2.8372e-07, 7.8101e-08, 5.6567e-07, 1.7264e-07, 2.4808e-07,\n",
      "          3.0732e-07, 1.4145e-08, 4.1648e-08, 8.1704e-07, 7.7041e-08,\n",
      "          4.4941e-08, 1.6102e-07, 3.0860e-09, 6.3934e-08, 5.4475e-08,\n",
      "          1.1126e-07, 1.7569e-07, 6.4398e-08, 1.3524e-06, 1.9022e-07,\n",
      "          1.4414e-08, 4.2404e-07, 3.2083e-05, 2.5686e-08]],\n",
      "\n",
      "        [[2.6942e-05, 1.2074e-04, 4.8584e-05, 2.5547e-05, 9.3548e-06,\n",
      "          5.6124e-05, 1.7595e-05, 1.2222e-04, 2.7384e-05, 6.6210e-05,\n",
      "          2.6982e-06, 2.1675e-05, 1.3014e-04, 4.7020e-05, 3.7890e-05,\n",
      "          1.0042e-05, 1.4822e-04, 3.8119e-05, 2.6237e-04, 2.1133e-06,\n",
      "          8.0704e-05, 1.1035e-05, 1.2069e-05, 5.6741e-06, 8.0070e-05,\n",
      "          1.0556e-06, 1.1963e-05, 5.9539e-05, 1.5508e-04, 4.8543e-05,\n",
      "          1.4957e-04, 1.2688e-04, 6.0693e-05, 1.6633e-03, 3.7387e-05,\n",
      "          1.8462e-05, 4.0941e-05, 6.4943e-05, 8.5201e-05, 7.6500e-06,\n",
      "          3.9154e-05, 8.5497e-06, 6.8973e-05, 2.3875e-05, 1.2301e-05,\n",
      "          3.8693e-05, 1.7959e-06, 5.7166e-06, 5.9762e-05, 1.0874e-05,\n",
      "          5.0593e-06, 2.2315e-05, 1.6171e-07, 2.8426e-06, 9.7221e-06,\n",
      "          2.6315e-05, 4.9587e-06, 8.3551e-06, 1.4186e-04, 2.3505e-05,\n",
      "          8.6113e-07, 2.4218e-05, 3.9752e-03, 6.6361e-06]]], device='cuda:0')}, 11: {'step': tensor(19527.), 'exp_avg': tensor([-0.0190,  0.0017,  0.0280, -0.1075, -0.0346, -0.0398,  0.0006, -0.0419,\n",
      "        -0.0479,  0.0154,  0.0038, -0.0598, -0.0199, -0.0066, -0.0564, -0.0386,\n",
      "        -0.0587,  0.0240,  0.0169, -0.0019, -0.0063, -0.0573, -0.0462,  0.0031,\n",
      "         0.0122, -0.0029, -0.0440, -0.0648,  0.0276,  0.0665, -0.0052, -0.0642,\n",
      "        -0.0045,  0.1899,  0.0572,  0.0215,  0.0069, -0.0294, -0.0284, -0.0050,\n",
      "         0.0049, -0.0750,  0.0352, -0.0083, -0.0191,  0.0255,  0.0149, -0.0098,\n",
      "         0.0210, -0.0347,  0.0262,  0.0198, -0.0045, -0.0056, -0.0355, -0.0290,\n",
      "        -0.0379, -0.0442,  0.0029, -0.0387, -0.0023, -0.0254,  0.2767, -0.0264],\n",
      "       device='cuda:0'), 'exp_avg_sq': tensor([1.3752e-01, 5.5958e-01, 2.7042e-03, 1.3816e+00, 2.1920e-01, 5.2651e-01,\n",
      "        2.8607e-01, 1.3146e+00, 5.6313e-01, 1.2801e-01, 1.5837e-01, 3.3253e-01,\n",
      "        5.2060e+00, 2.4990e-01, 7.2501e-02, 1.0154e-01, 1.2347e+01, 1.3754e-01,\n",
      "        1.0660e+00, 3.5343e-02, 2.7588e-01, 1.8591e-01, 1.0397e-01, 4.0335e-03,\n",
      "        3.1642e-01, 3.1536e-03, 2.9097e-01, 1.2154e+00, 3.4680e-01, 5.5511e-02,\n",
      "        9.6349e-01, 2.2709e+00, 5.6472e-01, 2.1724e+00, 4.2170e-02, 1.2840e-01,\n",
      "        6.0825e-02, 6.1321e-01, 1.7227e-01, 3.9809e-02, 2.5330e-01, 5.3122e-01,\n",
      "        1.2949e-01, 1.2572e-01, 8.1772e-01, 1.0989e-01, 2.7104e-02, 1.4578e-01,\n",
      "        6.3156e-01, 1.6898e-01, 1.7392e-01, 1.0098e-01, 5.6969e-03, 8.2159e-02,\n",
      "        1.0642e-01, 8.0035e-01, 1.2957e-01, 3.3427e-01, 8.1344e-01, 8.5503e-02,\n",
      "        3.1899e-02, 1.4359e+00, 5.4647e+00, 2.5687e-01], device='cuda:0')}, 12: {'step': tensor(19527.), 'exp_avg': tensor([[[[-4.2700e-03, -4.4195e-03, -3.7584e-03]],\n",
      "\n",
      "         [[-7.9508e-03, -8.4056e-03, -7.0962e-03]],\n",
      "\n",
      "         [[ 1.0322e-05, -4.6036e-06,  8.6828e-06]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-6.4900e-04, -6.2908e-04, -5.1265e-04]],\n",
      "\n",
      "         [[-1.0121e-03, -1.1361e-03, -9.0275e-04]],\n",
      "\n",
      "         [[-1.4162e-02, -1.4770e-02, -1.2545e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.1152e-04, -7.8127e-05,  1.3449e-04]],\n",
      "\n",
      "         [[-6.5503e-04, -6.2280e-04, -7.9981e-05]],\n",
      "\n",
      "         [[ 2.7753e-05,  3.7952e-05,  4.8451e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 2.4050e-04,  3.0226e-04,  3.1041e-04]],\n",
      "\n",
      "         [[-5.5860e-05, -3.9966e-05,  8.4953e-05]],\n",
      "\n",
      "         [[-9.5405e-04, -8.9911e-04, -8.2608e-05]]],\n",
      "\n",
      "\n",
      "        [[[ 2.6896e-03,  2.1818e-03,  2.8603e-03]],\n",
      "\n",
      "         [[ 4.3574e-03,  3.2980e-03,  4.9606e-03]],\n",
      "\n",
      "         [[-1.1598e-04, -1.5220e-04, -1.1897e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.4337e-04,  9.2425e-07,  9.0254e-05]],\n",
      "\n",
      "         [[-1.1669e-04, -3.3402e-04,  9.8172e-05]],\n",
      "\n",
      "         [[ 8.7172e-03,  7.0482e-03,  9.5332e-03]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 9.2603e-10, -5.3987e-08, -1.1761e-08]],\n",
      "\n",
      "         [[-1.2080e-08, -1.4462e-08, -6.7442e-08]],\n",
      "\n",
      "         [[-7.9702e-13, -2.1964e-12, -4.7594e-12]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-4.5179e-08, -4.4433e-08, -2.2793e-08]],\n",
      "\n",
      "         [[-1.2451e-08, -6.1697e-09, -2.4930e-09]],\n",
      "\n",
      "         [[-1.0058e-07, -5.2226e-08, -1.8447e-08]]],\n",
      "\n",
      "\n",
      "        [[[ 8.1239e-03,  8.0392e-03,  6.5221e-03]],\n",
      "\n",
      "         [[ 1.4308e-02,  1.4068e-02,  1.1476e-02]],\n",
      "\n",
      "         [[ 1.4474e-04,  1.3769e-04,  1.3787e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.6664e-03,  1.7023e-03,  1.4010e-03]],\n",
      "\n",
      "         [[ 2.0083e-03,  1.9251e-03,  1.6206e-03]],\n",
      "\n",
      "         [[ 2.5432e-02,  2.5107e-02,  2.0338e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.7989e-08, -2.1540e-08, -1.8733e-10]],\n",
      "\n",
      "         [[-5.6052e-45, -9.8063e-10, -1.0773e-09]],\n",
      "\n",
      "         [[ 5.6052e-45,  4.2039e-45, -5.6052e-45]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-2.1105e-08, -2.1101e-08, -3.9813e-22]],\n",
      "\n",
      "         [[ 5.6052e-45, -3.0733e-32, -3.4175e-32]],\n",
      "\n",
      "         [[-1.9009e-08, -2.9443e-08, -5.8423e-09]]]], device='cuda:0'), 'exp_avg_sq': tensor([[[[4.9924e-04, 6.0489e-04, 3.7458e-04]],\n",
      "\n",
      "         [[1.4055e-03, 1.7496e-03, 1.0620e-03]],\n",
      "\n",
      "         [[7.6132e-08, 8.0216e-08, 5.0948e-08]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[1.7724e-05, 1.7995e-05, 1.1952e-05]],\n",
      "\n",
      "         [[3.0186e-05, 3.5905e-05, 2.1727e-05]],\n",
      "\n",
      "         [[4.4032e-03, 5.4165e-03, 3.3237e-03]]],\n",
      "\n",
      "\n",
      "        [[[1.0319e-03, 1.0550e-03, 8.7602e-04]],\n",
      "\n",
      "         [[2.8385e-03, 2.8770e-03, 2.4007e-03]],\n",
      "\n",
      "         [[1.1076e-07, 1.0860e-07, 8.0544e-08]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[1.7149e-05, 1.7622e-05, 1.4112e-05]],\n",
      "\n",
      "         [[4.8973e-05, 4.8243e-05, 3.8016e-05]],\n",
      "\n",
      "         [[9.1301e-03, 9.2993e-03, 7.7476e-03]]],\n",
      "\n",
      "\n",
      "        [[[3.1652e-03, 3.6336e-03, 1.5993e-03]],\n",
      "\n",
      "         [[8.7212e-03, 1.0072e-02, 4.3222e-03]],\n",
      "\n",
      "         [[4.9874e-07, 5.1912e-07, 3.8301e-07]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[4.5670e-05, 4.9586e-05, 2.4217e-05]],\n",
      "\n",
      "         [[1.6705e-04, 1.8705e-04, 9.6027e-05]],\n",
      "\n",
      "         [[2.7964e-02, 3.2215e-02, 1.3949e-02]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[4.6385e-10, 4.9394e-08, 4.4326e-08]],\n",
      "\n",
      "         [[7.3170e-10, 1.1909e-07, 1.0858e-07]],\n",
      "\n",
      "         [[2.7369e-14, 4.1318e-14, 3.8789e-14]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[4.7415e-11, 1.0971e-09, 9.0104e-10]],\n",
      "\n",
      "         [[2.4928e-12, 1.2656e-09, 1.2016e-09]],\n",
      "\n",
      "         [[3.0270e-09, 3.8921e-07, 3.5178e-07]]],\n",
      "\n",
      "\n",
      "        [[[1.2064e-02, 1.1756e-02, 8.0123e-03]],\n",
      "\n",
      "         [[3.3218e-02, 3.2334e-02, 2.2055e-02]],\n",
      "\n",
      "         [[4.2716e-07, 4.0352e-07, 3.1580e-07]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[2.5531e-04, 2.4408e-04, 1.6268e-04]],\n",
      "\n",
      "         [[4.6524e-04, 4.4772e-04, 3.0828e-04]],\n",
      "\n",
      "         [[1.0733e-01, 1.0456e-01, 7.1272e-02]]],\n",
      "\n",
      "\n",
      "        [[[2.8842e-12, 3.0468e-12, 1.6285e-13]],\n",
      "\n",
      "         [[2.2689e-16, 2.5280e-13, 2.5285e-13]],\n",
      "\n",
      "         [[9.8070e-18, 3.8433e-17, 3.5540e-17]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.5770e-12, 5.5808e-12, 3.6556e-15]],\n",
      "\n",
      "         [[4.7107e-18, 1.2738e-15, 1.2784e-15]],\n",
      "\n",
      "         [[3.5381e-12, 4.3233e-12, 7.8633e-13]]]], device='cuda:0')}, 13: {'step': tensor(19527.), 'exp_avg': tensor([-2.4260e-02, -1.4979e-03,  1.1040e-02,  3.5452e-03,  9.3332e-03,\n",
      "         1.3147e-07, -7.9832e-02, -3.9549e-02, -1.9030e-02, -3.9518e-03,\n",
      "        -1.8712e-02, -4.5971e-05, -6.6013e-02, -1.7906e-02,  1.0424e-02,\n",
      "         4.4889e-02, -5.4116e-02,  5.0534e-07, -5.8689e-02,  1.6852e-08,\n",
      "         1.5680e-03, -1.8653e-02,  7.4902e-07, -4.1727e-02,  9.0926e-02,\n",
      "        -1.7705e-02,  4.7010e-02,  1.6861e-02, -3.6027e-02,  6.2314e-07,\n",
      "        -2.7411e-03, -6.0416e-02, -1.9576e-02,  2.0895e-02, -2.1531e-02,\n",
      "         2.9644e-02,  2.5995e-05,  9.7453e-02,  5.0200e-02,  9.5459e-04,\n",
      "        -2.2812e-02, -3.5980e-08, -9.7869e-03, -1.8127e-08, -1.2896e-09,\n",
      "         1.4183e-08,  4.3272e-07,  3.0948e-07,  9.4580e-04,  2.5524e-02,\n",
      "         1.3425e-02,  5.7789e-04,  8.4144e-02,  1.0818e-03, -8.1979e-02,\n",
      "        -1.1500e-02, -5.2033e-03,  1.3253e-01, -1.0777e-03,  3.1699e-03,\n",
      "        -1.6652e-02,  5.4717e-07,  4.1221e-02,  2.4079e-08], device='cuda:0'), 'exp_avg_sq': tensor([1.3869e-02, 2.3620e-02, 8.1803e-02, 2.0600e-02, 5.9817e-02, 1.5148e-11,\n",
      "        2.5742e-01, 3.3074e-02, 7.8078e-03, 1.1600e-03, 7.5033e-03, 9.5528e-06,\n",
      "        6.1241e-01, 1.1199e-02, 9.8383e-01, 8.3149e-02, 8.2156e-01, 3.4276e-10,\n",
      "        1.8353e-01, 3.8512e-06, 4.8750e-04, 7.6432e-03, 6.3455e-08, 1.2178e+00,\n",
      "        5.0281e-01, 6.4412e-03, 3.7279e-01, 4.4980e-02, 6.6730e-02, 3.5960e-09,\n",
      "        2.8075e-04, 6.6804e-01, 8.2233e-03, 6.5182e-02, 6.0227e-02, 5.1020e-02,\n",
      "        4.7271e-08, 9.2663e-01, 3.7544e+00, 6.7144e-02, 1.0631e+00, 2.4301e-04,\n",
      "        3.3738e-03, 5.6249e-11, 3.4722e-10, 4.6209e-13, 7.1909e-10, 2.8420e-10,\n",
      "        3.4387e-02, 2.7956e-01, 7.7943e-03, 7.0078e-04, 3.6423e-01, 2.5302e-05,\n",
      "        1.2350e+00, 3.9768e-03, 3.2803e-02, 1.6817e+00, 6.3927e-04, 1.1834e-03,\n",
      "        4.2473e-02, 9.4864e-07, 2.6655e-01, 5.0937e-12], device='cuda:0')}, 14: {'step': tensor(19527.), 'exp_avg': tensor([[[[-1.8606e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 6.1152e-03]]],\n",
      "\n",
      "\n",
      "        [[[-4.3735e-02]]],\n",
      "\n",
      "\n",
      "        [[[-2.9986e-02]]],\n",
      "\n",
      "\n",
      "        [[[-8.5535e-03]]],\n",
      "\n",
      "\n",
      "        [[[-3.2818e-08]]],\n",
      "\n",
      "\n",
      "        [[[ 1.0458e-01]]],\n",
      "\n",
      "\n",
      "        [[[-1.6057e-02]]],\n",
      "\n",
      "\n",
      "        [[[-3.8377e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 5.1743e-03]]],\n",
      "\n",
      "\n",
      "        [[[-4.0715e-03]]],\n",
      "\n",
      "\n",
      "        [[[-6.5531e-04]]],\n",
      "\n",
      "\n",
      "        [[[-2.5842e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 5.1151e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 5.3675e-02]]],\n",
      "\n",
      "\n",
      "        [[[-6.1159e-02]]],\n",
      "\n",
      "\n",
      "        [[[-9.2813e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 2.6824e-09]]],\n",
      "\n",
      "\n",
      "        [[[ 7.4607e-02]]],\n",
      "\n",
      "\n",
      "        [[[-2.5793e-08]]],\n",
      "\n",
      "\n",
      "        [[[-6.0157e-03]]],\n",
      "\n",
      "\n",
      "        [[[-3.3849e-03]]],\n",
      "\n",
      "\n",
      "        [[[-2.0639e-07]]],\n",
      "\n",
      "\n",
      "        [[[-3.0923e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 1.6491e-01]]],\n",
      "\n",
      "\n",
      "        [[[-6.3535e-03]]],\n",
      "\n",
      "\n",
      "        [[[-3.7845e-02]]],\n",
      "\n",
      "\n",
      "        [[[-2.7184e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 2.6861e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.6054e-07]]],\n",
      "\n",
      "\n",
      "        [[[-2.3680e-03]]],\n",
      "\n",
      "\n",
      "        [[[-7.4156e-02]]],\n",
      "\n",
      "\n",
      "        [[[-6.4318e-03]]],\n",
      "\n",
      "\n",
      "        [[[-8.0939e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 1.5585e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 5.7895e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.3404e-05]]],\n",
      "\n",
      "\n",
      "        [[[-2.7223e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 7.2648e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 9.5281e-04]]],\n",
      "\n",
      "\n",
      "        [[[ 6.1100e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 4.1467e-08]]],\n",
      "\n",
      "\n",
      "        [[[-9.8551e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 4.7704e-09]]],\n",
      "\n",
      "\n",
      "        [[[-2.6131e-08]]],\n",
      "\n",
      "\n",
      "        [[[ 2.0650e-09]]],\n",
      "\n",
      "\n",
      "        [[[ 2.0016e-07]]],\n",
      "\n",
      "\n",
      "        [[[-9.4584e-09]]],\n",
      "\n",
      "\n",
      "        [[[-1.6288e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.0093e-01]]],\n",
      "\n",
      "\n",
      "        [[[ 3.1312e-02]]],\n",
      "\n",
      "\n",
      "        [[[-4.0599e-03]]],\n",
      "\n",
      "\n",
      "        [[[-1.6607e-01]]],\n",
      "\n",
      "\n",
      "        [[[ 1.8946e-03]]],\n",
      "\n",
      "\n",
      "        [[[-4.1583e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.6889e-02]]],\n",
      "\n",
      "\n",
      "        [[[-1.3369e-02]]],\n",
      "\n",
      "\n",
      "        [[[-5.9712e-03]]],\n",
      "\n",
      "\n",
      "        [[[-4.0197e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 2.6509e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 6.8201e-03]]],\n",
      "\n",
      "\n",
      "        [[[-1.7965e-07]]],\n",
      "\n",
      "\n",
      "        [[[-9.5315e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 1.1530e-08]]]], device='cuda:0'), 'exp_avg_sq': tensor([[[[2.8185e-02]]],\n",
      "\n",
      "\n",
      "        [[[5.4749e-03]]],\n",
      "\n",
      "\n",
      "        [[[1.5727e-01]]],\n",
      "\n",
      "\n",
      "        [[[9.3650e-02]]],\n",
      "\n",
      "\n",
      "        [[[5.6502e-02]]],\n",
      "\n",
      "\n",
      "        [[[1.4235e-11]]],\n",
      "\n",
      "\n",
      "        [[[2.6531e-01]]],\n",
      "\n",
      "\n",
      "        [[[2.0782e-02]]],\n",
      "\n",
      "\n",
      "        [[[3.3769e-03]]],\n",
      "\n",
      "\n",
      "        [[[5.6570e-04]]],\n",
      "\n",
      "\n",
      "        [[[3.1185e-03]]],\n",
      "\n",
      "\n",
      "        [[[3.3245e-05]]],\n",
      "\n",
      "\n",
      "        [[[2.6113e-01]]],\n",
      "\n",
      "\n",
      "        [[[5.1760e-03]]],\n",
      "\n",
      "\n",
      "        [[[2.0101e-01]]],\n",
      "\n",
      "\n",
      "        [[[4.6625e-02]]],\n",
      "\n",
      "\n",
      "        [[[6.7892e-01]]],\n",
      "\n",
      "\n",
      "        [[[1.7130e-10]]],\n",
      "\n",
      "\n",
      "        [[[9.4835e-02]]],\n",
      "\n",
      "\n",
      "        [[[1.1632e-06]]],\n",
      "\n",
      "\n",
      "        [[[1.7953e-03]]],\n",
      "\n",
      "\n",
      "        [[[3.1775e-03]]],\n",
      "\n",
      "\n",
      "        [[[3.2208e-08]]],\n",
      "\n",
      "\n",
      "        [[[9.0200e-01]]],\n",
      "\n",
      "\n",
      "        [[[1.4890e+00]]],\n",
      "\n",
      "\n",
      "        [[[4.2534e-03]]],\n",
      "\n",
      "\n",
      "        [[[5.1140e-02]]],\n",
      "\n",
      "\n",
      "        [[[3.3975e-02]]],\n",
      "\n",
      "\n",
      "        [[[2.1467e-02]]],\n",
      "\n",
      "\n",
      "        [[[2.3603e-09]]],\n",
      "\n",
      "\n",
      "        [[[5.3060e-04]]],\n",
      "\n",
      "\n",
      "        [[[5.6201e-01]]],\n",
      "\n",
      "\n",
      "        [[[4.0452e-03]]],\n",
      "\n",
      "\n",
      "        [[[1.3488e-02]]],\n",
      "\n",
      "\n",
      "        [[[1.5871e-02]]],\n",
      "\n",
      "\n",
      "        [[[1.3079e-01]]],\n",
      "\n",
      "\n",
      "        [[[1.4480e-07]]],\n",
      "\n",
      "\n",
      "        [[[8.7157e-02]]],\n",
      "\n",
      "\n",
      "        [[[5.9819e-01]]],\n",
      "\n",
      "\n",
      "        [[[2.2435e-02]]],\n",
      "\n",
      "\n",
      "        [[[3.0105e-01]]],\n",
      "\n",
      "\n",
      "        [[[5.6771e-05]]],\n",
      "\n",
      "\n",
      "        [[[8.0476e-03]]],\n",
      "\n",
      "\n",
      "        [[[6.4415e-11]]],\n",
      "\n",
      "\n",
      "        [[[2.4633e-10]]],\n",
      "\n",
      "\n",
      "        [[[5.6294e-13]]],\n",
      "\n",
      "\n",
      "        [[[5.0993e-09]]],\n",
      "\n",
      "\n",
      "        [[[1.9245e-09]]],\n",
      "\n",
      "\n",
      "        [[[1.0556e-02]]],\n",
      "\n",
      "\n",
      "        [[[5.9387e-01]]],\n",
      "\n",
      "\n",
      "        [[[4.6131e-02]]],\n",
      "\n",
      "\n",
      "        [[[2.9238e-04]]],\n",
      "\n",
      "\n",
      "        [[[3.4601e-01]]],\n",
      "\n",
      "\n",
      "        [[[5.8278e-05]]],\n",
      "\n",
      "\n",
      "        [[[5.1629e-01]]],\n",
      "\n",
      "\n",
      "        [[[2.4037e-02]]],\n",
      "\n",
      "\n",
      "        [[[5.2045e-02]]],\n",
      "\n",
      "\n",
      "        [[[1.4849e-01]]],\n",
      "\n",
      "\n",
      "        [[[1.0194e-03]]],\n",
      "\n",
      "\n",
      "        [[[1.0632e-03]]],\n",
      "\n",
      "\n",
      "        [[[1.1249e-02]]],\n",
      "\n",
      "\n",
      "        [[[3.4961e-07]]],\n",
      "\n",
      "\n",
      "        [[[3.0626e-02]]],\n",
      "\n",
      "\n",
      "        [[[7.5592e-12]]]], device='cuda:0')}, 15: {'step': tensor(19527.), 'exp_avg': tensor([-2.4260e-02, -1.4985e-03,  1.1040e-02,  3.5450e-03,  9.3331e-03,\n",
      "         1.7528e-07, -7.9832e-02, -3.9549e-02, -1.9030e-02, -3.9523e-03,\n",
      "        -1.8712e-02, -4.6178e-05, -6.6013e-02, -1.7906e-02,  1.0425e-02,\n",
      "         4.4888e-02, -5.4116e-02,  7.0405e-07, -5.8689e-02,  6.7536e-07,\n",
      "         1.5682e-03, -1.8652e-02,  6.8834e-07, -4.1727e-02,  9.0926e-02,\n",
      "        -1.7705e-02,  4.7010e-02,  1.6861e-02, -3.6026e-02,  7.5475e-10,\n",
      "        -2.7413e-03, -6.0416e-02, -1.9576e-02,  2.0894e-02, -2.1531e-02,\n",
      "         2.9645e-02,  2.6059e-05,  9.7452e-02,  5.0199e-02,  9.5465e-04,\n",
      "        -2.2812e-02,  1.8342e-09, -9.7870e-03,  1.3782e-07, -3.1944e-09,\n",
      "         1.1581e-08,  2.5162e-07,  3.7808e-07,  9.4522e-04,  2.5524e-02,\n",
      "         1.3426e-02,  5.7826e-04,  8.4144e-02,  1.0818e-03, -8.1980e-02,\n",
      "        -1.1499e-02, -5.2029e-03,  1.3253e-01, -1.0769e-03,  3.1704e-03,\n",
      "        -1.6653e-02,  2.1683e-07,  4.1221e-02,  5.1985e-08], device='cuda:0'), 'exp_avg_sq': tensor([1.3869e-02, 2.3620e-02, 8.1803e-02, 2.0600e-02, 5.9817e-02, 1.5172e-11,\n",
      "        2.5742e-01, 3.3074e-02, 7.8078e-03, 1.1600e-03, 7.5033e-03, 9.5528e-06,\n",
      "        6.1241e-01, 1.1199e-02, 9.8383e-01, 8.3149e-02, 8.2156e-01, 3.4301e-10,\n",
      "        1.8353e-01, 3.8512e-06, 4.8750e-04, 7.6432e-03, 6.3455e-08, 1.2178e+00,\n",
      "        5.0281e-01, 6.4412e-03, 3.7279e-01, 4.4980e-02, 6.6730e-02, 3.5955e-09,\n",
      "        2.8075e-04, 6.6804e-01, 8.2233e-03, 6.5182e-02, 6.0227e-02, 5.1020e-02,\n",
      "        4.7271e-08, 9.2663e-01, 3.7544e+00, 6.7144e-02, 1.0631e+00, 2.4301e-04,\n",
      "        3.3738e-03, 5.6271e-11, 3.4722e-10, 4.6011e-13, 7.1896e-10, 2.8424e-10,\n",
      "        3.4387e-02, 2.7956e-01, 7.7943e-03, 7.0078e-04, 3.6423e-01, 2.5302e-05,\n",
      "        1.2350e+00, 3.9768e-03, 3.2803e-02, 1.6817e+00, 6.3927e-04, 1.1834e-03,\n",
      "        4.2473e-02, 9.4863e-07, 2.6655e-01, 5.1024e-12], device='cuda:0')}, 16: {'step': tensor(19527.), 'exp_avg': tensor([-0.0032,  0.0038,  0.0729,  0.0507, -0.0019,  0.0877, -0.0151, -0.0115,\n",
      "        -0.0021,  0.0092,  0.0138, -0.0104, -0.0276,  0.0085,  0.0728, -0.0100,\n",
      "         0.0197, -0.1011,  0.2166,  0.0211,  0.0356, -0.0384, -0.0711, -0.0024,\n",
      "        -0.1568,  0.0285, -0.0054,  0.0057,  0.1335, -0.0184,  0.0168, -0.0044,\n",
      "        -0.0523,  0.0030,  0.0236, -0.1453,  0.0305,  0.0196,  0.1018,  0.0049,\n",
      "        -0.0587, -0.1201,  0.0299, -0.1086,  0.0354, -0.1145, -0.0534, -0.1101,\n",
      "        -0.0105,  0.1273, -0.1175, -0.0096,  0.0205, -0.0584, -0.0654,  0.0347,\n",
      "         0.0265,  0.0823,  0.0150, -0.0584,  0.0644, -0.0662,  0.0082, -0.0172],\n",
      "       device='cuda:0'), 'exp_avg_sq': tensor([5.8268e-02, 3.4471e-02, 6.6755e-02, 5.9464e-02, 6.3880e-02, 3.0947e-01,\n",
      "        2.3148e+00, 6.8103e-03, 3.0688e-02, 6.3045e-01, 1.5407e-02, 1.3398e-01,\n",
      "        3.6026e-01, 1.7945e-02, 4.8302e-01, 2.9755e-02, 4.9727e-01, 6.5113e-01,\n",
      "        9.8654e+00, 1.8954e-01, 1.3266e-02, 5.8562e-02, 1.6722e+00, 6.4896e-01,\n",
      "        2.1557e+00, 3.1537e-02, 7.0228e-02, 3.7786e-02, 3.0213e+00, 2.1249e-02,\n",
      "        1.7774e-02, 1.1143e+00, 9.9601e-02, 8.2051e-03, 8.0365e-01, 1.1054e+00,\n",
      "        5.7623e-02, 2.4596e-01, 1.3864e+00, 8.0446e-02, 1.4233e-01, 3.4116e-01,\n",
      "        7.5049e-02, 2.7408e-01, 1.9930e-02, 2.9192e-01, 1.0220e-01, 3.8419e-01,\n",
      "        1.2200e-02, 2.3926e-01, 4.6962e-01, 5.7881e-03, 1.0680e-01, 6.9815e-02,\n",
      "        1.7681e+00, 8.2183e-02, 4.2900e-02, 6.3706e-01, 3.5049e-01, 5.0108e-01,\n",
      "        1.1024e+00, 2.0674e-01, 1.1213e-01, 7.2700e-02], device='cuda:0')}, 17: {'step': tensor(19527.), 'exp_avg': tensor([-0.0148,  0.1014,  0.0376,  0.0359,  0.0110, -0.0969, -0.0511, -0.0245,\n",
      "        -0.0012, -0.0160,  0.0174,  0.0183, -0.1000,  0.0025,  0.0356,  0.0672,\n",
      "        -0.0378,  0.1075, -0.3008, -0.0221,  0.0701, -0.0448,  0.0693, -0.0191,\n",
      "         0.2261,  0.0244,  0.0614,  0.0555, -0.2023,  0.0191, -0.0118, -0.0741,\n",
      "        -0.0342,  0.0785, -0.0708,  0.1709, -0.0325,  0.1602,  0.0677,  0.0591,\n",
      "         0.0088,  0.1256,  0.0239,  0.1198, -0.0372,  0.1202,  0.0573,  0.1156,\n",
      "         0.0399,  0.0535,  0.1359,  0.0432,  0.0750,  0.0629, -0.0727,  0.0171,\n",
      "         0.0259,  0.1678,  0.0103,  0.0631, -0.1048,  0.0644,  0.1138,  0.0157],\n",
      "       device='cuda:0'), 'exp_avg_sq': tensor([3.9959e-02, 5.3429e-01, 1.1061e-01, 6.0611e-02, 6.1305e-02, 3.4331e-01,\n",
      "        2.8551e+00, 1.4525e-02, 3.4422e-02, 6.4874e-01, 2.0057e-02, 1.5999e-01,\n",
      "        2.3462e+00, 3.6951e-02, 1.1644e+00, 1.5501e-01, 8.6096e-01, 7.0027e-01,\n",
      "        1.3596e+01, 2.1642e-01, 9.5263e-02, 7.1594e-02, 1.8909e+00, 8.4160e-01,\n",
      "        4.0719e+00, 3.3295e-02, 3.4101e-01, 1.6347e-01, 5.0271e+00, 2.4477e-02,\n",
      "        1.0830e-02, 2.0722e+00, 4.5625e-02, 2.5513e-01, 1.9456e+00, 1.4210e+00,\n",
      "        6.3949e-02, 1.7892e+00, 2.9640e+00, 4.5639e-01, 9.2821e-01, 4.0356e-01,\n",
      "        4.4507e-02, 3.3754e-01, 2.4219e-02, 3.4540e-01, 1.1584e-01, 4.3463e-01,\n",
      "        1.4074e-01, 4.0971e-01, 6.1581e-01, 7.4404e-02, 2.4364e-01, 8.2426e-02,\n",
      "        1.3608e+00, 2.7770e-02, 8.9922e-02, 2.0206e+00, 4.7188e-01, 5.5410e-01,\n",
      "        1.7875e+00, 2.2359e-01, 9.7265e-01, 8.3048e-02], device='cuda:0')}, 18: {'step': tensor(19527.), 'exp_avg': tensor([-5.8005e-08, -5.6541e-08, -5.7476e-08,  2.4781e-09, -1.1654e-08,\n",
      "        -5.8521e-08, -1.6558e-09, -1.0804e-08, -5.6669e-08,  4.2195e-09,\n",
      "         2.2039e-10,  1.0436e-08, -5.9690e-08, -1.3117e-09, -7.3251e-10,\n",
      "        -4.6287e-08, -6.0354e-08, -5.1124e-08,  3.0119e-09, -6.6554e-09,\n",
      "        -4.6129e-08, -1.5297e-08, -7.3659e-09, -2.6278e-09,  8.5631e-11,\n",
      "        -5.6943e-08,  4.9192e-10,  2.3016e-09, -9.9334e-09, -5.7441e-08,\n",
      "        -1.0003e-08,  8.8875e-09,  3.3455e-09, -2.4321e-09, -5.6321e-08,\n",
      "        -4.4654e-08,  1.8234e-09, -5.8855e-08, -5.6672e-08, -5.7641e-08,\n",
      "        -2.0762e-08, -1.7745e-08, -5.9186e-08, -1.5766e-08, -4.4791e-09,\n",
      "        -5.1634e-08,  9.0514e-10, -2.7833e-08, -5.4486e-08, -4.0561e-09,\n",
      "         7.7685e-10, -6.7034e-09, -5.5539e-08, -5.7603e-08,  1.0367e-09,\n",
      "        -5.4281e-09, -6.0148e-08, -6.5006e-08, -6.3870e-08,  1.4773e-09,\n",
      "        -5.2516e-09, -8.8818e-09, -1.9415e-08, -5.1032e-08, -1.1720e-08,\n",
      "        -4.6724e-08,  9.4396e-10, -5.9035e-08, -7.6803e-09, -1.8038e-08,\n",
      "         2.9524e-09,  7.4653e-09, -5.9495e-08, -5.4057e-08, -3.4788e-08,\n",
      "        -3.9277e-09, -5.0918e-10, -5.5766e-08,  4.7422e-08, -5.7990e-08,\n",
      "        -1.1003e-08, -5.6637e-08, -1.8707e-08, -4.3261e-08, -5.9017e-08,\n",
      "        -8.1377e-09, -6.2545e-09, -3.3969e-08,  5.7515e-08, -4.3736e-09,\n",
      "        -5.6984e-08, -5.5975e-08, -6.2317e-08, -5.5294e-08, -4.3451e-09,\n",
      "        -6.0750e-08, -6.1929e-08, -1.7000e-08,  1.5613e-09, -5.8126e-08,\n",
      "         1.8494e-09, -1.0112e-08, -5.7210e-08, -5.8516e-08,  6.1358e-10,\n",
      "         1.4744e-09, -5.8255e-08, -1.7437e-09,  1.2891e-08, -4.1786e-10,\n",
      "         4.2373e-08, -1.5112e-08, -2.3173e-09, -1.3561e-08, -5.7612e-08,\n",
      "         1.8607e-09,  2.0266e-09, -6.2197e-08, -6.7615e-08, -6.5089e-08,\n",
      "        -1.4971e-08, -2.2823e-09, -5.9983e-08, -4.6677e-08, -5.9215e-08,\n",
      "        -5.7091e-08, -4.3897e-08, -6.0818e-08, -2.9406e-10, -1.0992e-08,\n",
      "        -1.9473e-08, -1.9632e-09, -1.2239e-09, -3.2506e-09, -2.2226e-08,\n",
      "         1.3387e-09,  2.0677e-09,  1.3162e-08, -5.5938e-08, -1.1231e-09,\n",
      "        -1.1626e-09, -1.8691e-10, -1.7418e-08, -1.1285e-08, -9.2094e-12,\n",
      "        -4.6753e-09, -5.9298e-08, -5.5726e-08, -3.2728e-09,  1.5421e-09,\n",
      "         1.4163e-09, -5.7893e-08, -1.1904e-08, -2.1228e-10, -6.0949e-08,\n",
      "        -1.4903e-08, -7.8257e-09, -4.4613e-08, -5.6923e-08, -5.8518e-08,\n",
      "        -1.0249e-08, -5.7144e-08, -5.9825e-09, -5.7923e-08, -4.8167e-11,\n",
      "        -3.8094e-08, -3.9661e-08, -5.6235e-08, -3.8042e-08,  1.3928e-09,\n",
      "         1.9934e-09, -5.6102e-08, -7.0630e-12, -5.4614e-08, -5.8663e-08,\n",
      "         2.0695e-09, -5.6791e-10,  3.1158e-09,  7.5422e-10,  7.4172e-10,\n",
      "        -9.1107e-12,  2.6334e-11, -2.8045e-08, -9.2094e-12, -1.0280e-08,\n",
      "        -2.1238e-09, -1.4531e-08, -2.6398e-08, -7.3332e-09, -7.4966e-12,\n",
      "        -5.9558e-09,  4.0440e-12,  2.7587e-11,  2.4421e-11,  1.3180e-11,\n",
      "         1.5695e-11, -1.0245e-09, -8.0386e-12, -9.2112e-12, -8.8018e-12,\n",
      "        -9.1926e-12, -5.2336e-08, -5.7027e-08, -6.0485e-08, -5.8657e-08,\n",
      "        -2.4333e-10, -1.0703e-08, -6.7196e-08, -6.0189e-08, -1.3994e-08,\n",
      "        -5.7941e-08, -5.7820e-08, -1.1979e-09,  4.2768e-10,  9.7114e-09,\n",
      "         2.4500e-09,  3.3110e-09,  9.6575e-10,  5.3109e-09, -8.7383e-10,\n",
      "        -4.6908e-08, -5.7077e-08, -5.6859e-08, -4.5867e-09, -1.5671e-08,\n",
      "         1.0478e-10, -6.5905e-09, -5.3617e-08, -6.1889e-10, -1.5097e-08,\n",
      "        -1.6809e-08,  8.2058e-10,  2.3254e-09, -5.2185e-08, -1.9076e-08,\n",
      "        -5.7930e-08, -5.9280e-08, -3.3753e-08, -1.8647e-10, -6.3003e-08,\n",
      "         3.0023e-09, -5.9008e-08, -4.2611e-08, -1.1948e-08,  6.3990e-10,\n",
      "        -5.3133e-08, -5.6951e-09,  2.0953e-09, -2.2774e-08, -2.4831e-09,\n",
      "        -5.3853e-08, -5.7761e-08, -5.8817e-08, -1.8116e-08,  2.6916e-09,\n",
      "        -2.3597e-09, -5.4761e-08, -9.0603e-10, -2.1789e-08, -5.3554e-08,\n",
      "        -5.8189e-08,  1.4027e-09, -1.2238e-08, -5.7746e-08, -2.0875e-09,\n",
      "        -4.8821e-08, -1.9290e-09, -5.6387e-09, -2.0260e-09, -1.7095e-09,\n",
      "         8.4660e-10, -6.0867e-08, -5.6222e-08, -2.5379e-09, -1.5011e-08,\n",
      "        -1.2298e-09, -9.2633e-09,  3.1719e-09,  2.8508e-09,  4.7206e-09,\n",
      "         1.8103e-09, -4.7991e-09, -4.7208e-09, -2.6368e-09, -5.7746e-08,\n",
      "        -5.7507e-08, -5.5645e-08, -1.7834e-08, -3.9811e-09,  1.8550e-09,\n",
      "        -1.5613e-09, -5.8469e-08, -5.6094e-08, -2.0653e-08, -5.7250e-08,\n",
      "         2.4900e-09, -4.5148e-08, -4.9065e-08, -9.0694e-09,  1.7996e-08,\n",
      "        -5.2546e-08, -5.7078e-08, -1.6836e-08,  4.4030e-09, -5.0037e-08,\n",
      "         5.4468e-10, -4.6352e-08], device='cuda:0'), 'exp_avg_sq': tensor([3.7190e-06, 9.2586e-06, 5.9679e-05, 4.6636e-06, 2.3888e-06, 4.6054e-06,\n",
      "        6.6496e-06, 7.1898e-07, 1.4591e-05, 3.4965e-06, 9.0936e-06, 1.0094e-05,\n",
      "        3.6340e-05, 6.1289e-06, 9.3932e-06, 5.0650e-05, 1.5406e-06, 2.1268e-05,\n",
      "        9.0216e-06, 7.9235e-07, 1.1883e-06, 2.6715e-06, 4.9917e-06, 5.2862e-06,\n",
      "        4.1813e-06, 1.3692e-05, 2.7976e-06, 7.1107e-06, 1.3090e-06, 4.4109e-06,\n",
      "        1.1476e-06, 1.0781e-05, 2.0979e-06, 2.7121e-06, 3.5542e-05, 1.6644e-06,\n",
      "        6.0192e-06, 1.0245e-05, 5.3970e-06, 7.1381e-06, 1.8756e-07, 6.9657e-05,\n",
      "        3.3256e-06, 1.4021e-07, 3.8014e-06, 2.7827e-05, 8.0423e-06, 7.2523e-07,\n",
      "        6.1219e-05, 3.7527e-06, 3.3268e-06, 6.6741e-07, 1.1060e-05, 2.4053e-05,\n",
      "        9.6648e-06, 1.9296e-06, 9.4072e-06, 1.5022e-06, 2.2786e-05, 3.7543e-06,\n",
      "        6.2343e-06, 3.0011e-07, 2.4658e-06, 2.6475e-05, 2.0703e-06, 6.0118e-06,\n",
      "        6.8897e-06, 5.0613e-06, 4.6933e-06, 2.2977e-05, 1.1357e-05, 1.0604e-06,\n",
      "        1.0402e-05, 1.0168e-06, 2.0253e-06, 5.7846e-05, 2.0563e-06, 3.2136e-05,\n",
      "        1.9678e-06, 5.2485e-05, 2.4997e-07, 7.0281e-05, 1.2947e-05, 2.1317e-06,\n",
      "        8.5273e-06, 8.4679e-07, 3.9898e-06, 1.9506e-06, 1.5257e-06, 4.1086e-06,\n",
      "        4.7952e-05, 4.1590e-05, 8.0330e-06, 4.6231e-06, 6.1792e-07, 7.1663e-06,\n",
      "        2.6282e-06, 5.1624e-07, 7.2817e-06, 6.4409e-06, 3.2340e-06, 4.1814e-05,\n",
      "        2.6354e-05, 5.3495e-06, 7.2484e-06, 7.0178e-06, 4.9461e-05, 9.5414e-06,\n",
      "        7.3271e-05, 5.7733e-05, 1.3422e-07, 3.0672e-06, 6.1882e-06, 2.9777e-06,\n",
      "        6.7419e-05, 1.5582e-06, 9.8022e-06, 5.1355e-07, 2.3435e-07, 1.2846e-06,\n",
      "        2.1514e-07, 9.9618e-07, 5.9565e-07, 1.3386e-05, 2.3358e-05, 1.4278e-05,\n",
      "        3.0993e-06, 5.0418e-07, 1.2742e-05, 4.0033e-06, 3.0265e-06, 8.8298e-06,\n",
      "        1.0432e-05, 4.5987e-06, 1.2396e-06, 6.8715e-06, 1.0541e-05, 5.6603e-06,\n",
      "        3.7820e-06, 5.7778e-06, 1.9195e-06, 3.5132e-07, 1.3165e-05, 2.0869e-06,\n",
      "        7.6484e-06, 5.4376e-06, 3.7419e-05, 6.0006e-06, 6.4886e-06, 5.1213e-06,\n",
      "        5.1217e-06, 1.9923e-05, 1.7709e-06, 1.5514e-06, 2.7444e-05, 3.2150e-06,\n",
      "        4.9573e-06, 2.3027e-06, 1.5070e-05, 3.5199e-07, 2.2198e-06, 1.9893e-05,\n",
      "        6.3727e-07, 1.3063e-05, 3.1325e-06, 1.7396e-06, 2.0842e-06, 1.1382e-05,\n",
      "        1.1234e-06, 3.0187e-06, 3.0148e-06, 1.0014e-05, 7.6484e-06, 3.8973e-05,\n",
      "        3.1245e-06, 5.0319e-06, 1.5512e-06, 3.0119e-06, 3.0826e-06, 3.0738e-06,\n",
      "        7.6484e-06, 7.6484e-06, 2.0333e-06, 7.6484e-06, 2.6405e-06, 4.8075e-06,\n",
      "        1.7433e-06, 6.2973e-07, 2.7788e-06, 7.6484e-06, 4.8796e-06, 7.6484e-06,\n",
      "        7.6484e-06, 7.6484e-06, 7.6484e-06, 7.6484e-06, 1.0287e-05, 7.6484e-06,\n",
      "        7.6484e-06, 7.6484e-06, 7.6484e-06, 3.7031e-05, 1.9297e-05, 1.3280e-07,\n",
      "        6.0707e-05, 8.0427e-06, 3.6624e-06, 1.0896e-07, 2.3276e-07, 6.6772e-05,\n",
      "        2.9415e-06, 7.9987e-06, 7.3549e-06, 3.5515e-06, 5.1572e-06, 6.6102e-06,\n",
      "        5.3341e-06, 5.1320e-06, 5.7598e-06, 4.9306e-06, 4.7654e-05, 2.1324e-05,\n",
      "        4.5670e-05, 3.5708e-06, 3.0986e-06, 1.1447e-05, 6.6606e-06, 5.3575e-06,\n",
      "        7.8881e-06, 2.6636e-06, 1.8004e-06, 7.7083e-06, 8.8152e-06, 1.4054e-06,\n",
      "        6.0443e-06, 5.9488e-06, 1.1258e-05, 6.0232e-07, 9.9967e-07, 5.1807e-06,\n",
      "        2.9155e-06, 4.3426e-06, 1.9108e-06, 2.6176e-06, 4.6656e-06, 6.6074e-05,\n",
      "        1.0835e-06, 1.0699e-05, 2.4260e-06, 6.6572e-06, 2.3289e-06, 3.0766e-05,\n",
      "        1.5154e-05, 2.6757e-05, 4.5956e-06, 5.1400e-06, 2.2466e-05, 1.0352e-05,\n",
      "        2.4820e-06, 8.8805e-07, 5.0638e-06, 3.0157e-06, 4.9509e-07, 7.5746e-06,\n",
      "        1.2418e-06, 1.1845e-07, 1.3096e-05, 1.8000e-06, 1.6429e-06, 4.0129e-06,\n",
      "        2.3169e-06, 2.6807e-05, 7.0444e-05, 7.7718e-06, 4.0744e-06, 9.7007e-06,\n",
      "        4.3612e-06, 1.1012e-05, 4.3173e-06, 4.3272e-06, 6.2587e-06, 2.7636e-06,\n",
      "        4.6167e-06, 3.9163e-06, 5.5357e-05, 5.5106e-05, 2.0235e-05, 5.4404e-07,\n",
      "        1.9510e-06, 7.9488e-06, 2.4207e-06, 2.1499e-05, 2.0363e-05, 6.1376e-07,\n",
      "        5.9514e-06, 3.0278e-06, 1.3315e-06, 2.7469e-07, 7.6056e-05, 2.8644e-06,\n",
      "        2.1895e-05, 1.4696e-05, 1.5157e-06, 1.0539e-05, 4.2686e-07, 6.2050e-06,\n",
      "        1.5291e-06], device='cuda:0')}, 19: {'step': tensor(19527.), 'exp_avg': tensor([[-1.8147e-08, -1.8779e-08, -1.9317e-08,  ..., -1.8294e-08,\n",
      "         -1.9784e-08, -1.9717e-08],\n",
      "        [-3.7364e-08, -3.9051e-08, -6.5671e-08,  ..., -4.6091e-08,\n",
      "         -6.1292e-08, -4.7471e-08],\n",
      "        [-1.2372e-08, -1.2664e-08, -1.2262e-08,  ..., -1.2036e-08,\n",
      "         -1.2436e-08, -1.2644e-08],\n",
      "        ...,\n",
      "        [ 3.7588e-08,  3.8419e-08,  4.2867e-08,  ...,  4.0382e-08,\n",
      "          4.2589e-08,  4.1383e-08],\n",
      "        [-7.8276e-08, -8.0523e-08, -7.9322e-08,  ..., -8.0653e-08,\n",
      "         -8.0521e-08, -8.1466e-08],\n",
      "        [-1.8443e-08, -1.9104e-08, -2.1488e-08,  ..., -1.9684e-08,\n",
      "         -2.1998e-08, -2.1070e-08]], device='cuda:0'), 'exp_avg_sq': tensor([[1.1838e-07, 1.1774e-07, 9.2655e-08,  ..., 1.1745e-07, 1.1275e-07,\n",
      "         1.1758e-07],\n",
      "        [2.3430e-09, 2.3304e-09, 1.8339e-09,  ..., 2.3244e-09, 2.2318e-09,\n",
      "         2.3274e-09],\n",
      "        [1.8699e-07, 1.8598e-07, 1.4642e-07,  ..., 1.8550e-07, 1.7810e-07,\n",
      "         1.8572e-07],\n",
      "        ...,\n",
      "        [5.7037e-09, 5.6730e-09, 4.4644e-09,  ..., 5.6591e-09, 5.4333e-09,\n",
      "         5.6658e-09],\n",
      "        [3.0710e-08, 3.0546e-08, 2.4045e-08,  ..., 3.0469e-08, 2.9260e-08,\n",
      "         3.0511e-08],\n",
      "        [7.1770e-10, 7.1377e-10, 5.6162e-10,  ..., 7.1198e-10, 6.8347e-10,\n",
      "         7.1274e-10]], device='cuda:0')}, 20: {'step': tensor(19527.), 'exp_avg': tensor([ 1.5683e-07, -7.5325e-09,  3.7747e-08,  4.1139e-08, -2.3845e-08,\n",
      "         5.3073e-09, -1.7556e-09,  4.3223e-08,  2.0762e-10,  6.7836e-09,\n",
      "         2.6033e-08,  2.6696e-08,  1.6122e-07,  8.5206e-08,  1.9161e-07,\n",
      "         2.1700e-08,  4.6083e-08,  1.8957e-07, -1.4869e-08,  5.9916e-08,\n",
      "        -1.3025e-09,  1.0395e-07,  5.2849e-08,  3.5022e-08, -5.1388e-09,\n",
      "         9.0990e-09, -1.6800e-08, -6.1807e-09, -1.9238e-10,  2.5830e-08,\n",
      "         1.3547e-08,  3.9193e-08,  3.0726e-08,  1.7021e-07,  4.3953e-08,\n",
      "         8.5355e-08,  1.8382e-07, -8.0776e-09,  4.1138e-09, -1.0982e-08,\n",
      "         6.6019e-08,  6.6989e-09,  1.4378e-10,  4.9362e-09,  3.6367e-09,\n",
      "         1.2491e-08,  8.9565e-08,  6.9917e-09,  6.9564e-08,  2.8892e-07,\n",
      "        -4.2759e-09,  1.3465e-08, -3.4284e-09, -2.8534e-08,  2.5579e-07,\n",
      "         1.3132e-07,  4.0589e-08, -1.5877e-08,  4.9584e-08,  7.4740e-08,\n",
      "        -1.4020e-09,  5.1222e-09, -7.1171e-09,  5.3279e-08], device='cuda:0'), 'exp_avg_sq': tensor([2.6348e-09, 9.4920e-10, 2.5813e-08, 1.3941e-08, 7.8332e-09, 1.3382e-11,\n",
      "        5.2982e-09, 1.0596e-08, 2.4123e-11, 2.3779e-10, 3.4249e-11, 2.5665e-11,\n",
      "        1.7810e-08, 5.6233e-10, 2.6449e-08, 4.7367e-08, 3.1104e-08, 6.3151e-11,\n",
      "        2.8770e-10, 1.5617e-10, 3.5678e-10, 6.5032e-11, 2.2375e-10, 1.5248e-07,\n",
      "        2.8824e-11, 2.7448e-11, 9.0920e-08, 8.3848e-09, 1.3102e-10, 5.6278e-12,\n",
      "        8.1290e-12, 1.1015e-08, 2.0630e-11, 1.0419e-08, 3.1040e-10, 1.7862e-10,\n",
      "        2.7599e-11, 5.6914e-08, 7.9005e-08, 5.1830e-09, 1.5980e-07, 1.8592e-10,\n",
      "        6.7805e-11, 1.5920e-11, 2.9366e-11, 3.6069e-10, 1.1891e-11, 2.8914e-10,\n",
      "        8.2494e-09, 1.6533e-08, 1.8073e-10, 1.4213e-09, 2.3832e-07, 1.8494e-09,\n",
      "        8.9024e-08, 6.0509e-09, 2.3183e-08, 6.7610e-08, 2.0875e-11, 3.7327e-11,\n",
      "        2.5826e-10, 1.6126e-10, 2.2685e-08, 6.8839e-11], device='cuda:0')}, 21: {'step': tensor(19527.), 'exp_avg': tensor([[[-5.6649e-06,  1.5122e-04, -2.1817e-05,  1.6790e-03, -3.8945e-04,\n",
      "           2.4821e-04,  3.1909e-05,  1.6874e-06, -2.7045e-05, -1.7740e-04,\n",
      "           8.9981e-05,  5.3536e-06],\n",
      "         [-1.4818e-05,  7.8443e-05, -9.4827e-07,  2.0546e-04, -5.1693e-04,\n",
      "           1.2317e-04,  3.4712e-05, -3.1960e-05,  3.2940e-06,  2.0568e-05,\n",
      "           2.4277e-05,  1.2641e-05],\n",
      "         [-5.9341e-05, -8.1877e-05,  6.2751e-06, -2.3324e-04, -7.0246e-05,\n",
      "          -6.6211e-05, -1.1276e-04,  5.9246e-05,  1.1712e-05,  1.2855e-04,\n",
      "          -2.4254e-05, -6.7837e-06],\n",
      "         [ 2.3370e-04, -1.3800e-03, -9.2333e-05, -6.0787e-04,  1.1580e-03,\n",
      "           1.1241e-03,  2.0432e-03, -6.1587e-04, -1.9752e-04, -1.2552e-03,\n",
      "           1.0392e-03, -2.8867e-04],\n",
      "         [-3.9599e-06,  6.5370e-04, -5.6036e-06,  1.0209e-03, -1.8917e-04,\n",
      "          -9.9900e-05, -3.4030e-04,  3.2716e-04,  6.3834e-05,  3.2436e-04,\n",
      "          -2.6723e-04,  8.0634e-05],\n",
      "         [-4.3615e-06,  4.3804e-04,  3.8278e-05, -2.0251e-04,  1.0363e-04,\n",
      "          -2.8448e-04, -5.7643e-04,  2.3366e-04,  8.4625e-05,  2.8918e-04,\n",
      "          -2.0368e-04,  5.7680e-05],\n",
      "         [-8.0866e-05, -8.5381e-04, -2.7573e-05, -2.5783e-03,  3.9471e-04,\n",
      "           3.1294e-05,  1.7312e-04, -1.2806e-04, -1.7926e-05,  1.1541e-05,\n",
      "           5.2829e-05, -4.3245e-05],\n",
      "         [ 6.0229e-05,  2.9438e-04, -4.9276e-06,  2.6090e-03,  2.6515e-04,\n",
      "           1.0340e-04,  5.8647e-05, -3.0139e-05, -4.0108e-05, -1.3377e-04,\n",
      "           1.2691e-04,  2.7310e-05],\n",
      "         [ 9.9081e-06, -9.4726e-05, -1.7074e-06, -4.7955e-04,  6.7553e-05,\n",
      "           1.8329e-05,  4.0039e-04, -2.0106e-04, -2.7702e-06,  5.4255e-05,\n",
      "           3.4451e-05, -2.4915e-05],\n",
      "         [-3.1567e-05, -1.4363e-04, -2.7928e-05, -8.5719e-04,  5.2919e-04,\n",
      "          -1.8924e-04, -1.3541e-04,  1.3549e-04, -1.4946e-05,  5.0361e-05,\n",
      "          -1.3432e-04, -5.6024e-05],\n",
      "         [-1.5470e-04, -2.1241e-04,  2.3747e-05, -2.1504e-03,  2.6006e-04,\n",
      "          -2.2517e-04, -9.3428e-05,  3.5528e-04,  5.9264e-05,  4.1498e-04,\n",
      "          -1.7766e-04,  6.3547e-05],\n",
      "         [-8.2003e-05, -1.4094e-03, -6.6750e-05, -3.9804e-03,  1.3815e-04,\n",
      "           1.6564e-05,  2.2640e-04, -2.4987e-04, -6.6031e-05, -1.9664e-04,\n",
      "           2.8191e-04, -4.9805e-05]]], device='cuda:0'), 'exp_avg_sq': tensor([[[4.8028e-08, 6.7265e-06, 8.1653e-08, 2.5748e-06, 2.3777e-06,\n",
      "          1.4828e-06, 1.2768e-05, 1.3139e-06, 2.2542e-07, 4.5518e-06,\n",
      "          5.4020e-07, 8.6719e-08],\n",
      "         [3.4402e-08, 8.2782e-07, 2.1755e-08, 5.0515e-07, 7.4603e-07,\n",
      "          8.3416e-07, 4.4509e-06, 7.0510e-07, 5.3260e-08, 6.5335e-07,\n",
      "          4.0546e-07, 4.0098e-08],\n",
      "         [4.8945e-08, 3.8600e-06, 2.6811e-08, 1.0864e-06, 1.5450e-06,\n",
      "          1.4763e-06, 1.4524e-05, 9.4593e-07, 1.6860e-07, 3.1105e-06,\n",
      "          1.4489e-06, 7.0276e-08],\n",
      "         [6.1493e-08, 2.5438e-06, 4.7079e-08, 5.2811e-07, 1.5248e-06,\n",
      "          2.0469e-06, 7.3055e-06, 5.5779e-07, 6.8285e-08, 2.3198e-06,\n",
      "          1.0831e-06, 1.4208e-07],\n",
      "         [8.8089e-08, 8.3658e-06, 1.2716e-07, 2.2592e-06, 1.6638e-06,\n",
      "          2.4293e-06, 1.6855e-05, 3.4012e-06, 3.5845e-07, 7.6642e-06,\n",
      "          2.1985e-06, 2.7092e-07],\n",
      "         [4.4300e-08, 3.8876e-06, 1.1681e-07, 1.2568e-06, 1.2837e-06,\n",
      "          1.4355e-06, 5.3769e-06, 6.8925e-07, 1.4538e-07, 1.7868e-06,\n",
      "          5.3206e-07, 7.8997e-08],\n",
      "         [1.6208e-07, 1.7587e-05, 3.6221e-07, 6.4554e-06, 3.1392e-06,\n",
      "          2.4809e-06, 9.8419e-06, 9.4958e-06, 3.5611e-07, 1.7409e-05,\n",
      "          1.4066e-06, 1.4225e-07],\n",
      "         [4.9905e-07, 1.4254e-05, 3.0896e-07, 5.8125e-06, 9.6020e-06,\n",
      "          1.3668e-05, 3.4784e-05, 5.3035e-06, 1.4423e-06, 1.6404e-05,\n",
      "          6.7025e-06, 4.1753e-07],\n",
      "         [1.3096e-07, 2.5609e-06, 1.6823e-07, 1.2811e-06, 4.3423e-06,\n",
      "          3.5893e-06, 2.7662e-05, 4.3197e-06, 1.5258e-07, 1.2281e-05,\n",
      "          1.4278e-06, 5.2276e-07],\n",
      "         [1.8692e-07, 9.6999e-06, 3.4442e-07, 2.9439e-06, 7.8891e-06,\n",
      "          9.0613e-06, 1.6581e-05, 8.2380e-06, 4.7714e-07, 4.9442e-06,\n",
      "          4.9225e-06, 6.8493e-07],\n",
      "         [1.4311e-07, 2.4363e-05, 6.5800e-07, 5.8900e-06, 4.5557e-06,\n",
      "          6.7754e-06, 1.5292e-05, 8.5763e-06, 7.0515e-07, 6.9667e-06,\n",
      "          2.1814e-06, 7.0955e-07],\n",
      "         [3.6412e-07, 4.1286e-05, 8.6094e-07, 1.1463e-05, 3.9869e-06,\n",
      "          4.6770e-06, 1.7601e-05, 7.3446e-06, 6.7880e-07, 2.0950e-05,\n",
      "          8.0431e-06, 1.2497e-07]]], device='cuda:0')}, 22: {'step': tensor(19527.), 'exp_avg': tensor([[ 2.3182e-05,  1.5900e-04,  1.9279e-04,  1.8541e-04,  6.9345e-05,\n",
      "          2.4214e-05,  1.3367e-04,  7.7413e-05,  1.8111e-04,  1.5344e-04,\n",
      "          1.2655e-04,  7.8797e-05],\n",
      "        [ 4.3943e-03,  4.9012e-03,  6.0739e-03,  9.4245e-03,  3.8386e-03,\n",
      "          4.9413e-03,  5.7990e-03,  4.0079e-03,  6.5423e-03,  5.8858e-03,\n",
      "          5.1210e-03,  2.2810e-03],\n",
      "        [ 7.3335e-03,  8.3457e-03,  9.3414e-03,  1.5090e-02,  5.8522e-03,\n",
      "          7.2699e-03,  9.4765e-03,  6.2511e-03,  1.0787e-02,  1.0165e-02,\n",
      "          6.5533e-03,  7.4992e-03],\n",
      "        [-2.9562e-03, -2.1456e-03, -1.1571e-03,  4.2775e-03, -5.6211e-03,\n",
      "         -3.2175e-03, -2.0882e-03, -3.1158e-03, -3.7205e-04, -1.8887e-03,\n",
      "         -3.4534e-03, -4.9415e-03],\n",
      "        [-4.4941e-03, -3.3780e-03, -2.3302e-03,  1.4619e-03, -5.7882e-03,\n",
      "         -3.6309e-03, -4.8944e-03, -4.3455e-03, -2.1674e-03, -2.5936e-03,\n",
      "         -3.3688e-03, -4.4920e-03],\n",
      "        [ 6.5302e-03,  7.5109e-03,  8.1536e-03,  1.3048e-02,  6.6900e-03,\n",
      "          4.7846e-03,  6.1105e-03,  6.1720e-03,  9.3688e-03,  8.6335e-03,\n",
      "          6.7676e-03,  6.1320e-03],\n",
      "        [-9.2363e-03, -8.4554e-03, -7.6958e-03, -5.1822e-03, -6.3147e-03,\n",
      "         -7.8754e-03, -6.6340e-03, -8.1644e-03, -8.8965e-03, -8.9037e-03,\n",
      "         -5.8082e-03, -7.4589e-03],\n",
      "        [ 1.8593e-03,  1.4359e-03,  1.5996e-03,  1.2411e-02,  5.5274e-03,\n",
      "         -7.4617e-04,  8.5141e-04, -1.5573e-03,  1.8157e-03,  2.5499e-03,\n",
      "          1.4211e-04,  1.5258e-03],\n",
      "        [-1.1373e-02, -1.1427e-02, -1.1771e-02, -9.0646e-03, -1.3858e-02,\n",
      "         -1.1775e-02, -1.3110e-02, -1.3109e-02, -1.0522e-02, -1.1413e-02,\n",
      "         -1.1193e-02, -1.1410e-02],\n",
      "        [ 1.0665e-02,  1.0749e-02,  1.1136e-02,  1.2847e-02,  6.6425e-03,\n",
      "          1.3377e-02,  1.2547e-02,  1.3816e-02,  1.0713e-02,  1.2221e-02,\n",
      "          1.3108e-02,  1.1606e-02],\n",
      "        [ 6.0987e-03,  3.9531e-03,  1.6687e-03, -1.3273e-02,  9.6399e-03,\n",
      "          4.2961e-03,  5.2999e-03,  3.2211e-03,  1.4885e-04,  2.4342e-03,\n",
      "          5.1298e-04,  3.7776e-03],\n",
      "        [-8.8441e-03, -1.1649e-02, -1.5212e-02, -4.1225e-02, -6.6764e-03,\n",
      "         -7.4485e-03, -1.3491e-02, -3.2528e-03, -1.7597e-02, -1.7242e-02,\n",
      "         -8.5071e-03, -4.5994e-03]], device='cuda:0'), 'exp_avg_sq': tensor([[1.2911e-06, 8.0512e-07, 7.6660e-07, 1.8567e-06, 6.1940e-07, 1.0269e-06,\n",
      "         1.4580e-06, 5.9622e-07, 1.8279e-06, 1.7065e-06, 8.7137e-07, 7.6170e-07],\n",
      "        [6.7898e-04, 6.4829e-04, 6.7941e-04, 7.6268e-04, 6.5551e-04, 6.4648e-04,\n",
      "         6.3844e-04, 6.2853e-04, 8.1750e-04, 7.8969e-04, 6.3776e-04, 5.3155e-04],\n",
      "        [1.4838e-03, 1.4133e-03, 1.2527e-03, 1.5667e-03, 1.3799e-03, 1.3913e-03,\n",
      "         1.3655e-03, 1.3721e-03, 1.6179e-03, 1.6408e-03, 1.2481e-03, 1.3334e-03],\n",
      "        [1.5545e-03, 1.5446e-03, 1.5422e-03, 2.0022e-03, 1.4973e-03, 1.4600e-03,\n",
      "         1.4346e-03, 1.4775e-03, 1.7428e-03, 1.7702e-03, 1.4628e-03, 1.4119e-03],\n",
      "        [1.9374e-03, 1.8228e-03, 1.7981e-03, 1.9852e-03, 1.5574e-03, 1.8109e-03,\n",
      "         1.6033e-03, 1.8152e-03, 2.1368e-03, 2.1720e-03, 1.7078e-03, 1.6729e-03],\n",
      "        [1.9115e-03, 1.6980e-03, 1.6499e-03, 1.8051e-03, 1.6852e-03, 1.3668e-03,\n",
      "         1.5289e-03, 1.6754e-03, 2.0417e-03, 2.0306e-03, 1.4982e-03, 1.5732e-03],\n",
      "        [2.6552e-03, 2.5933e-03, 2.5231e-03, 2.7255e-03, 2.5917e-03, 2.4190e-03,\n",
      "         1.9233e-03, 2.3712e-03, 2.9445e-03, 3.3152e-03, 2.1488e-03, 2.3431e-03],\n",
      "        [2.4992e-03, 2.4065e-03, 2.3718e-03, 2.5802e-03, 2.3673e-03, 2.3144e-03,\n",
      "         2.1076e-03, 2.4615e-03, 2.6916e-03, 2.7073e-03, 2.2380e-03, 2.3323e-03],\n",
      "        [2.7879e-03, 2.7022e-03, 2.6615e-03, 2.9409e-03, 2.6247e-03, 2.6449e-03,\n",
      "         2.6409e-03, 2.6562e-03, 4.6150e-03, 3.1495e-03, 2.5940e-03, 2.8080e-03],\n",
      "        [3.7864e-03, 3.6373e-03, 3.6288e-03, 3.8976e-03, 3.6414e-03, 3.5385e-03,\n",
      "         3.3238e-03, 3.6029e-03, 4.1264e-03, 5.7399e-03, 3.5221e-03, 3.7036e-03],\n",
      "        [8.7434e-03, 8.4217e-03, 8.2982e-03, 8.9257e-03, 8.4383e-03, 8.1437e-03,\n",
      "         7.7312e-03, 8.2961e-03, 9.4111e-03, 9.5019e-03, 8.9998e-03, 7.8727e-03],\n",
      "        [1.6038e-02, 1.5562e-02, 1.5376e-02, 1.6713e-02, 1.5477e-02, 1.5212e-02,\n",
      "         1.4704e-02, 1.5449e-02, 1.7236e-02, 1.7632e-02, 1.5111e-02, 1.5853e-02]],\n",
      "       device='cuda:0')}, 23: {'step': tensor(19527.), 'exp_avg': tensor([3.5916, 3.6000, 3.5948, 3.6335, 3.6129, 3.6013, 3.5985, 3.5736, 3.5926,\n",
      "        3.5799, 3.6170, 3.6101], device='cuda:0'), 'exp_avg_sq': tensor([251.3727, 252.1559, 252.0064, 252.4582, 251.3412, 251.0986, 250.8487,\n",
      "        252.1084, 252.5350, 252.5894, 250.9989, 251.8431], device='cuda:0')}, 24: {'step': tensor(19527.), 'exp_avg': tensor([[-3.0384e-02, -3.0354e-02, -3.0386e-02, -2.9961e-02, -3.0548e-02,\n",
      "         -3.0401e-02, -3.0346e-02, -3.0330e-02, -3.0371e-02, -3.0666e-02,\n",
      "         -3.0472e-02, -3.0407e-02],\n",
      "        [-9.3598e-04, -9.3758e-04, -9.4687e-04, -9.1942e-04, -9.5689e-04,\n",
      "         -9.3490e-04, -9.2547e-04, -9.4723e-04, -9.4918e-04, -1.0033e-03,\n",
      "         -9.3564e-04, -9.3130e-04],\n",
      "        [-5.2506e-02, -5.2480e-02, -5.2500e-02, -5.1949e-02, -5.2710e-02,\n",
      "         -5.2525e-02, -5.2438e-02, -5.2411e-02, -5.2475e-02, -5.2805e-02,\n",
      "         -5.2626e-02, -5.2552e-02],\n",
      "        [-3.3088e-02, -3.3061e-02, -3.3080e-02, -3.2664e-02, -3.3240e-02,\n",
      "         -3.3103e-02, -3.3054e-02, -3.3020e-02, -3.3060e-02, -3.3309e-02,\n",
      "         -3.3179e-02, -3.3115e-02],\n",
      "        [ 9.5229e-03,  9.4991e-03,  9.4871e-03,  9.2849e-03,  9.6011e-03,\n",
      "          9.5139e-03,  9.5442e-03,  9.4678e-03,  9.4612e-03,  9.5099e-03,\n",
      "          9.5749e-03,  9.5160e-03],\n",
      "        [-1.0841e-02, -1.0834e-02, -1.0847e-02, -1.0705e-02, -1.0898e-02,\n",
      "         -1.0846e-02, -1.0821e-02, -1.0828e-02, -1.0844e-02, -1.0961e-02,\n",
      "         -1.0869e-02, -1.0848e-02],\n",
      "        [ 5.8094e-03,  5.8253e-03,  5.8622e-03,  5.7690e-03,  5.8779e-03,\n",
      "          5.8079e-03,  5.7863e-03,  5.8552e-03,  5.8659e-03,  5.9700e-03,\n",
      "          5.8020e-03,  5.8260e-03],\n",
      "        [-3.5536e-02, -3.5504e-02, -3.5545e-02, -3.5041e-02, -3.5734e-02,\n",
      "         -3.5558e-02, -3.5486e-02, -3.5480e-02, -3.5531e-02, -3.5906e-02,\n",
      "         -3.5640e-02, -3.5565e-02],\n",
      "        [-8.9281e-03, -8.9222e-03, -8.9334e-03, -8.8156e-03, -8.9760e-03,\n",
      "         -8.9323e-03, -8.9112e-03, -8.9175e-03, -8.9308e-03, -9.0282e-03,\n",
      "         -8.9505e-03, -8.9339e-03],\n",
      "        [ 5.3465e-03,  5.3507e-03,  5.3637e-03,  5.2727e-03,  5.3993e-03,\n",
      "          5.3439e-03,  5.3285e-03,  5.3583e-03,  5.3641e-03,  5.4567e-03,\n",
      "          5.3591e-03,  5.3522e-03],\n",
      "        [-1.0191e-02, -1.0184e-02, -1.0197e-02, -1.0063e-02, -1.0246e-02,\n",
      "         -1.0196e-02, -1.0172e-02, -1.0179e-02, -1.0194e-02, -1.0305e-02,\n",
      "         -1.0217e-02, -1.0198e-02],\n",
      "        [-1.6414e-02, -1.6401e-02, -1.6419e-02, -1.6211e-02, -1.6493e-02,\n",
      "         -1.6423e-02, -1.6386e-02, -1.6389e-02, -1.6414e-02, -1.6581e-02,\n",
      "         -1.6455e-02, -1.6424e-02],\n",
      "        [-2.6543e-02, -2.6512e-02, -2.6517e-02, -2.6193e-02, -2.6644e-02,\n",
      "         -2.6561e-02, -2.6529e-02, -2.6465e-02, -2.6501e-02, -2.6641e-02,\n",
      "         -2.6629e-02, -2.6578e-02],\n",
      "        [-8.2113e-03, -8.2065e-03, -8.2244e-03, -8.0913e-03, -8.2730e-03,\n",
      "         -8.2144e-03, -8.1930e-03, -8.2114e-03, -8.2228e-03, -8.3530e-03,\n",
      "         -8.2317e-03, -8.2123e-03],\n",
      "        [ 1.6025e-02,  1.6038e-02,  1.5986e-02,  1.6078e-02,  1.5948e-02,\n",
      "          1.6028e-02,  1.6002e-02,  1.5958e-02,  1.5983e-02,  1.5774e-02,\n",
      "          1.6038e-02,  1.6071e-02],\n",
      "        [ 9.1318e-03,  9.0632e-03,  9.0577e-03,  8.7630e-03,  9.1589e-03,\n",
      "          9.1604e-03,  9.2265e-03,  9.0064e-03,  9.0057e-03,  9.0313e-03,\n",
      "          9.1935e-03,  9.1102e-03],\n",
      "        [-5.7571e-02, -5.7556e-02, -5.7569e-02, -5.7037e-02, -5.7770e-02,\n",
      "         -5.7588e-02, -5.7485e-02, -5.7474e-02, -5.7546e-02, -5.7856e-02,\n",
      "         -5.7692e-02, -5.7632e-02],\n",
      "        [-2.2427e-02, -2.2412e-02, -2.2440e-02, -2.2145e-02, -2.2548e-02,\n",
      "         -2.2438e-02, -2.2385e-02, -2.2401e-02, -2.2434e-02, -2.2677e-02,\n",
      "         -2.2484e-02, -2.2442e-02],\n",
      "        [ 1.9592e-03,  1.9711e-03,  1.9858e-03,  1.9853e-03,  1.9901e-03,\n",
      "          1.9606e-03,  1.9458e-03,  1.9955e-03,  1.9957e-03,  2.0484e-03,\n",
      "          1.9529e-03,  1.9594e-03],\n",
      "        [-3.0177e-03, -3.0101e-03, -3.0081e-03, -2.9905e-03, -3.0099e-03,\n",
      "         -3.0238e-03, -3.0194e-03, -2.9989e-03, -3.0055e-03, -3.0060e-03,\n",
      "         -3.0245e-03, -3.0184e-03],\n",
      "        [-4.2859e-03, -4.2876e-03, -4.2956e-03, -4.2389e-03, -4.3159e-03,\n",
      "         -4.2842e-03, -4.2732e-03, -4.2906e-03, -4.2949e-03, -4.3538e-03,\n",
      "         -4.2922e-03, -4.2869e-03],\n",
      "        [-2.0838e-02, -2.0824e-02, -2.0849e-02, -2.0576e-02, -2.0948e-02,\n",
      "         -2.0848e-02, -2.0799e-02, -2.0813e-02, -2.0844e-02, -2.1069e-02,\n",
      "         -2.0891e-02, -2.0852e-02],\n",
      "        [-1.4385e-03, -1.4309e-03, -1.4261e-03, -1.4322e-03, -1.4192e-03,\n",
      "         -1.4446e-03, -1.4442e-03, -1.4194e-03, -1.4240e-03, -1.4025e-03,\n",
      "         -1.4412e-03, -1.4379e-03],\n",
      "        [-8.3370e-02, -8.3258e-02, -8.3296e-02, -8.2154e-02, -8.3745e-02,\n",
      "         -8.3427e-02, -8.3333e-02, -8.3129e-02, -8.3236e-02, -8.3796e-02,\n",
      "         -8.3650e-02, -8.3460e-02],\n",
      "        [ 7.5937e-05,  8.2557e-05,  8.6763e-05,  6.6294e-05,  1.0390e-04,\n",
      "          6.9494e-05,  6.8220e-05,  9.3714e-05,  8.8553e-05,  1.1861e-04,\n",
      "          7.9349e-05,  7.6440e-05],\n",
      "        [-1.7898e-02, -1.7886e-02, -1.7908e-02, -1.7671e-02, -1.7995e-02,\n",
      "         -1.7907e-02, -1.7866e-02, -1.7877e-02, -1.7903e-02, -1.8096e-02,\n",
      "         -1.7944e-02, -1.7911e-02],\n",
      "        [ 2.1309e-03,  2.0727e-03,  1.9964e-03,  2.0464e-03,  1.9843e-03,\n",
      "          2.1586e-03,  2.2344e-03,  1.9581e-03,  1.9602e-03,  1.5556e-03,\n",
      "          2.1821e-03,  2.1623e-03],\n",
      "        [-1.4677e-02, -1.4676e-02, -1.4722e-02, -1.4454e-02, -1.4816e-02,\n",
      "         -1.4678e-02, -1.4633e-02, -1.4703e-02, -1.4721e-02, -1.5032e-02,\n",
      "         -1.4707e-02, -1.4672e-02],\n",
      "        [ 3.4192e-03,  3.4253e-03,  3.4401e-03,  3.4091e-03,  3.4529e-03,\n",
      "          3.4198e-03,  3.4055e-03,  3.4428e-03,  3.4469e-03,  3.5073e-03,\n",
      "          3.4180e-03,  3.4217e-03],\n",
      "        [-3.6517e-03, -3.6477e-03, -3.6507e-03, -3.6080e-03, -3.6656e-03,\n",
      "         -3.6544e-03, -3.6466e-03, -3.6435e-03, -3.6494e-03, -3.6805e-03,\n",
      "         -3.6609e-03, -3.6539e-03],\n",
      "        [-8.7958e-03, -8.7903e-03, -8.8002e-03, -8.6806e-03, -8.8426e-03,\n",
      "         -8.8005e-03, -8.7803e-03, -8.7843e-03, -8.7980e-03, -8.8937e-03,\n",
      "         -8.8202e-03, -8.8045e-03],\n",
      "        [-3.1067e-02, -3.1052e-02, -3.1057e-02, -3.0757e-02, -3.1178e-02,\n",
      "         -3.1079e-02, -3.1020e-02, -3.1005e-02, -3.1048e-02, -3.1206e-02,\n",
      "         -3.1145e-02, -3.1105e-02],\n",
      "        [-1.1402e-02, -1.1394e-02, -1.1408e-02, -1.1259e-02, -1.1462e-02,\n",
      "         -1.1407e-02, -1.1380e-02, -1.1388e-02, -1.1405e-02, -1.1530e-02,\n",
      "         -1.1430e-02, -1.1409e-02],\n",
      "        [-1.3650e-02, -1.3657e-02, -1.3698e-02, -1.3488e-02, -1.3772e-02,\n",
      "         -1.3649e-02, -1.3596e-02, -1.3683e-02, -1.3703e-02, -1.3988e-02,\n",
      "         -1.3669e-02, -1.3645e-02],\n",
      "        [-6.4860e-03, -6.4702e-03, -6.4667e-03, -6.3828e-03, -6.4840e-03,\n",
      "         -6.4954e-03, -6.4925e-03, -6.4405e-03, -6.4512e-03, -6.4825e-03,\n",
      "         -6.5050e-03, -6.4832e-03],\n",
      "        [ 1.9141e-02,  1.9134e-02,  1.9161e-02,  1.8872e-02,  1.9272e-02,\n",
      "          1.9142e-02,  1.9101e-02,  1.9134e-02,  1.9153e-02,  1.9398e-02,\n",
      "          1.9194e-02,  1.9151e-02],\n",
      "        [-1.8009e-02, -1.7996e-02, -1.8016e-02, -1.7787e-02, -1.8098e-02,\n",
      "         -1.8020e-02, -1.7978e-02, -1.7983e-02, -1.8011e-02, -1.8196e-02,\n",
      "         -1.8055e-02, -1.8022e-02],\n",
      "        [ 4.0189e-03,  3.9860e-03,  3.9311e-03,  3.9741e-03,  3.9185e-03,\n",
      "          4.0356e-03,  4.0795e-03,  3.9033e-03,  3.9070e-03,  3.6267e-03,\n",
      "          4.0504e-03,  4.0450e-03],\n",
      "        [ 3.1111e-02,  3.1093e-02,  3.1028e-02,  3.0960e-02,  3.1049e-02,\n",
      "          3.1129e-02,  3.1106e-02,  3.0960e-02,  3.1002e-02,  3.0788e-02,\n",
      "          3.1170e-02,  3.1170e-02],\n",
      "        [-1.0803e-02, -1.0809e-02, -1.0839e-02, -1.0683e-02, -1.0894e-02,\n",
      "         -1.0802e-02, -1.0757e-02, -1.0828e-02, -1.0845e-02, -1.1057e-02,\n",
      "         -1.0820e-02, -1.0804e-02],\n",
      "        [-5.7508e-02, -5.7404e-02, -5.7383e-02, -5.6912e-02, -5.7518e-02,\n",
      "         -5.7598e-02, -5.7539e-02, -5.7241e-02, -5.7332e-02, -5.7314e-02,\n",
      "         -5.7658e-02, -5.7586e-02],\n",
      "        [ 1.2779e-03,  1.2829e-03,  1.2903e-03,  1.2517e-03,  1.3089e-03,\n",
      "          1.2741e-03,  1.2681e-03,  1.2920e-03,  1.2917e-03,  1.3385e-03,\n",
      "          1.2821e-03,  1.2804e-03],\n",
      "        [-2.2158e-02, -2.2142e-02, -2.2174e-02, -2.1846e-02, -2.2298e-02,\n",
      "         -2.2168e-02, -2.2121e-02, -2.2137e-02, -2.2167e-02, -2.2435e-02,\n",
      "         -2.2220e-02, -2.2173e-02],\n",
      "        [-1.1689e-02, -1.1681e-02, -1.1695e-02, -1.1542e-02, -1.1752e-02,\n",
      "         -1.1694e-02, -1.1667e-02, -1.1675e-02, -1.1692e-02, -1.1819e-02,\n",
      "         -1.1718e-02, -1.1697e-02],\n",
      "        [-9.6869e-03, -9.6820e-03, -9.6949e-03, -9.5640e-03, -9.7427e-03,\n",
      "         -9.6910e-03, -9.6676e-03, -9.6783e-03, -9.6924e-03, -9.8035e-03,\n",
      "         -9.7117e-03, -9.6939e-03],\n",
      "        [ 1.4750e-02,  1.4747e-02,  1.4772e-02,  1.4554e-02,  1.4856e-02,\n",
      "          1.4752e-02,  1.4714e-02,  1.4751e-02,  1.4770e-02,  1.4968e-02,\n",
      "          1.4789e-02,  1.4763e-02],\n",
      "        [-1.0116e-02, -1.0107e-02, -1.0118e-02, -9.9920e-03, -1.0163e-02,\n",
      "         -1.0122e-02, -1.0099e-02, -1.0099e-02, -1.0114e-02, -1.0213e-02,\n",
      "         -1.0142e-02, -1.0122e-02],\n",
      "        [ 1.2812e-03,  1.2876e-03,  1.2963e-03,  1.2524e-03,  1.3177e-03,\n",
      "          1.2759e-03,  1.2698e-03,  1.2990e-03,  1.2982e-03,  1.3529e-03,\n",
      "          1.2860e-03,  1.2840e-03],\n",
      "        [ 4.1913e-03,  4.1715e-03,  4.1553e-03,  4.1129e-03,  4.1726e-03,\n",
      "          4.2020e-03,  4.2166e-03,  4.1380e-03,  4.1434e-03,  4.0768e-03,\n",
      "          4.2152e-03,  4.2018e-03],\n",
      "        [-5.5769e-02, -5.5754e-02, -5.5775e-02, -5.5233e-02, -5.5984e-02,\n",
      "         -5.5785e-02, -5.5671e-02, -5.5686e-02, -5.5759e-02, -5.6111e-02,\n",
      "         -5.5891e-02, -5.5827e-02],\n",
      "        [ 2.8540e-03,  2.8559e-03,  2.8631e-03,  2.8012e-03,  2.8897e-03,\n",
      "          2.8506e-03,  2.8442e-03,  2.8617e-03,  2.8628e-03,  2.9223e-03,\n",
      "          2.8636e-03,  2.8554e-03],\n",
      "        [-5.2426e-03, -5.2391e-03, -5.2350e-03, -5.1983e-03, -5.2516e-03,\n",
      "         -5.2447e-03, -5.2368e-03, -5.2270e-03, -5.2348e-03, -5.2299e-03,\n",
      "         -5.2589e-03, -5.2555e-03],\n",
      "        [-2.6789e-02, -2.6921e-02, -2.6999e-02, -2.7153e-02, -2.6952e-02,\n",
      "         -2.6733e-02, -2.6482e-02, -2.7017e-02, -2.7069e-02, -2.7437e-02,\n",
      "         -2.6705e-02, -2.6798e-02],\n",
      "        [ 8.7227e-03,  8.7198e-03,  8.7337e-03,  8.5941e-03,  8.7837e-03,\n",
      "          8.7215e-03,  8.7035e-03,  8.7208e-03,  8.7293e-03,  8.8517e-03,\n",
      "          8.7446e-03,  8.7253e-03],\n",
      "        [-7.7780e-02, -7.7692e-02, -7.7729e-02, -7.6739e-02, -7.8117e-02,\n",
      "         -7.7828e-02, -7.7721e-02, -7.7579e-02, -7.7681e-02, -7.8203e-02,\n",
      "         -7.8018e-02, -7.7858e-02],\n",
      "        [-4.0084e-02, -4.0043e-02, -4.0089e-02, -3.9497e-02, -4.0316e-02,\n",
      "         -4.0108e-02, -4.0039e-02, -4.0014e-02, -4.0069e-02, -4.0496e-02,\n",
      "         -4.0204e-02, -4.0110e-02],\n",
      "        [-4.0173e-02, -4.0132e-02, -4.0157e-02, -3.9626e-02, -4.0358e-02,\n",
      "         -4.0197e-02, -4.0137e-02, -4.0081e-02, -4.0134e-02, -4.0447e-02,\n",
      "         -4.0294e-02, -4.0213e-02],\n",
      "        [ 5.0245e-03,  4.9877e-03,  4.9222e-03,  4.9830e-03,  4.9028e-03,\n",
      "          5.0442e-03,  5.0907e-03,  4.8890e-03,  4.8963e-03,  4.5647e-03,\n",
      "          5.0608e-03,  5.0565e-03],\n",
      "        [-1.4663e-02, -1.4651e-02, -1.4670e-02, -1.4465e-02, -1.4743e-02,\n",
      "         -1.4671e-02, -1.4637e-02, -1.4642e-02, -1.4665e-02, -1.4826e-02,\n",
      "         -1.4705e-02, -1.4677e-02],\n",
      "        [-8.1246e-03, -8.1158e-03, -8.1251e-03, -8.0253e-03, -8.1571e-03,\n",
      "         -8.1314e-03, -8.1124e-03, -8.1076e-03, -8.1219e-03, -8.1992e-03,\n",
      "         -8.1442e-03, -8.1301e-03],\n",
      "        [-3.7415e-04, -3.3773e-04, -3.3588e-04, -2.1406e-04, -3.7602e-04,\n",
      "         -3.8080e-04, -3.9407e-04, -3.0738e-04, -3.1869e-04, -3.0071e-04,\n",
      "         -4.2056e-04, -3.8797e-04],\n",
      "        [ 6.0105e-03,  6.0118e-03,  6.0239e-03,  5.9268e-03,  6.0631e-03,\n",
      "          6.0092e-03,  5.9929e-03,  6.0168e-03,  6.0240e-03,  6.1172e-03,\n",
      "          6.0267e-03,  6.0158e-03],\n",
      "        [-3.1230e-03, -3.1450e-03, -3.1880e-03, -3.1013e-03, -3.2168e-03,\n",
      "         -3.1129e-03, -3.0666e-03, -3.1977e-03, -3.2041e-03, -3.4465e-03,\n",
      "         -3.1123e-03, -3.1112e-03],\n",
      "        [-2.1041e-03, -2.0988e-03, -2.0976e-03, -2.0848e-03, -2.0986e-03,\n",
      "         -2.1083e-03, -2.1054e-03, -2.0913e-03, -2.0957e-03, -2.0958e-03,\n",
      "         -2.1089e-03, -2.1045e-03]], device='cuda:0'), 'exp_avg_sq': tensor([[1.2399e-01, 1.2449e-01, 1.2434e-01, 1.2454e-01, 1.2393e-01, 1.2380e-01,\n",
      "         1.2372e-01, 1.2460e-01, 1.2484e-01, 1.2476e-01, 1.2378e-01, 1.2448e-01],\n",
      "        [4.6857e-03, 4.7026e-03, 4.6979e-03, 4.7054e-03, 4.6832e-03, 4.6785e-03,\n",
      "         4.6742e-03, 4.7075e-03, 4.7156e-03, 4.7134e-03, 4.6777e-03, 4.7033e-03],\n",
      "        [1.4980e-01, 1.5041e-01, 1.5023e-01, 1.5048e-01, 1.4972e-01, 1.4957e-01,\n",
      "         1.4949e-01, 1.5054e-01, 1.5084e-01, 1.5072e-01, 1.4955e-01, 1.5042e-01],\n",
      "        [9.2402e-02, 9.2770e-02, 9.2661e-02, 9.2812e-02, 9.2350e-02, 9.2258e-02,\n",
      "         9.2202e-02, 9.2849e-02, 9.3033e-02, 9.2969e-02, 9.2244e-02, 9.2771e-02],\n",
      "        [9.1549e-03, 9.1904e-03, 9.1791e-03, 9.1953e-03, 9.1469e-03, 9.1392e-03,\n",
      "         9.1339e-03, 9.1977e-03, 9.2140e-03, 9.2113e-03, 9.1398e-03, 9.1928e-03],\n",
      "        [1.8586e-02, 1.8659e-02, 1.8637e-02, 1.8668e-02, 1.8576e-02, 1.8557e-02,\n",
      "         1.8544e-02, 1.8676e-02, 1.8712e-02, 1.8699e-02, 1.8553e-02, 1.8658e-02],\n",
      "        [1.3034e-02, 1.3083e-02, 1.3070e-02, 1.3091e-02, 1.3030e-02, 1.3016e-02,\n",
      "         1.3003e-02, 1.3092e-02, 1.3118e-02, 1.3117e-02, 1.3011e-02, 1.3077e-02],\n",
      "        [1.9160e-01, 1.9236e-01, 1.9213e-01, 1.9245e-01, 1.9149e-01, 1.9130e-01,\n",
      "         1.9117e-01, 1.9253e-01, 1.9291e-01, 1.9278e-01, 1.9127e-01, 1.9236e-01],\n",
      "        [1.3356e-02, 1.3408e-02, 1.3393e-02, 1.3415e-02, 1.3349e-02, 1.3336e-02,\n",
      "         1.3326e-02, 1.3420e-02, 1.3446e-02, 1.3437e-02, 1.3333e-02, 1.3407e-02],\n",
      "        [1.1598e-02, 1.1642e-02, 1.1629e-02, 1.1648e-02, 1.1593e-02, 1.1581e-02,\n",
      "         1.1571e-02, 1.1651e-02, 1.1673e-02, 1.1667e-02, 1.1577e-02, 1.1638e-02],\n",
      "        [1.7178e-02, 1.7245e-02, 1.7226e-02, 1.7254e-02, 1.7170e-02, 1.7152e-02,\n",
      "         1.7140e-02, 1.7261e-02, 1.7294e-02, 1.7283e-02, 1.7149e-02, 1.7245e-02],\n",
      "        [3.7302e-02, 3.7450e-02, 3.7406e-02, 3.7467e-02, 3.7282e-02, 3.7244e-02,\n",
      "         3.7218e-02, 3.7485e-02, 3.7557e-02, 3.7532e-02, 3.7237e-02, 3.7450e-02],\n",
      "        [3.3231e-02, 3.3369e-02, 3.3329e-02, 3.3383e-02, 3.3212e-02, 3.3178e-02,\n",
      "         3.3162e-02, 3.3397e-02, 3.3468e-02, 3.3446e-02, 3.3173e-02, 3.3364e-02],\n",
      "        [2.3039e-02, 2.3128e-02, 2.3102e-02, 2.3140e-02, 2.3027e-02, 2.3004e-02,\n",
      "         2.2986e-02, 2.3150e-02, 2.3194e-02, 2.3179e-02, 2.2999e-02, 2.3129e-02],\n",
      "        [7.3786e-02, 7.4047e-02, 7.3973e-02, 7.4093e-02, 7.3743e-02, 7.3661e-02,\n",
      "         7.3594e-02, 7.4129e-02, 7.4252e-02, 7.4229e-02, 7.3653e-02, 7.4047e-02],\n",
      "        [1.7621e-02, 1.7682e-02, 1.7670e-02, 1.7699e-02, 1.7610e-02, 1.7596e-02,\n",
      "         1.7565e-02, 1.7696e-02, 1.7727e-02, 1.7745e-02, 1.7597e-02, 1.7688e-02],\n",
      "        [1.4094e-01, 1.4151e-01, 1.4135e-01, 1.4157e-01, 1.4087e-01, 1.4073e-01,\n",
      "         1.4065e-01, 1.4162e-01, 1.4191e-01, 1.4180e-01, 1.4070e-01, 1.4150e-01],\n",
      "        [8.0094e-02, 8.0408e-02, 8.0316e-02, 8.0447e-02, 8.0052e-02, 7.9971e-02,\n",
      "         7.9912e-02, 8.0482e-02, 8.0637e-02, 8.0583e-02, 7.9953e-02, 8.0404e-02],\n",
      "        [4.7737e-03, 4.7915e-03, 4.7871e-03, 4.7945e-03, 4.7714e-03, 4.7666e-03,\n",
      "         4.7615e-03, 4.7962e-03, 4.8049e-03, 4.8033e-03, 4.7658e-03, 4.7930e-03],\n",
      "        [2.1995e-04, 2.2077e-04, 2.2063e-04, 2.2102e-04, 2.1975e-04, 2.1944e-04,\n",
      "         2.1941e-04, 2.2123e-04, 2.2169e-04, 2.2153e-04, 2.1959e-04, 2.2109e-04],\n",
      "        [4.1726e-03, 4.1883e-03, 4.1839e-03, 4.1906e-03, 4.1706e-03, 4.1666e-03,\n",
      "         4.1630e-03, 4.1921e-03, 4.1998e-03, 4.1976e-03, 4.1654e-03, 4.1874e-03],\n",
      "        [6.9494e-02, 6.9766e-02, 6.9687e-02, 6.9800e-02, 6.9458e-02, 6.9388e-02,\n",
      "         6.9337e-02, 6.9831e-02, 6.9965e-02, 6.9919e-02, 6.9372e-02, 6.9763e-02],\n",
      "        [1.0037e-03, 1.0071e-03, 1.0064e-03, 1.0080e-03, 1.0033e-03, 1.0023e-03,\n",
      "         1.0013e-03, 1.0079e-03, 1.0099e-03, 1.0097e-03, 1.0019e-03, 1.0065e-03],\n",
      "        [4.9032e-01, 4.9233e-01, 4.9173e-01, 4.9254e-01, 4.9003e-01, 4.8953e-01,\n",
      "         4.8925e-01, 4.9276e-01, 4.9377e-01, 4.9344e-01, 4.8946e-01, 4.9230e-01],\n",
      "        [1.9265e-03, 1.9335e-03, 1.9319e-03, 1.9349e-03, 1.9258e-03, 1.9239e-03,\n",
      "         1.9219e-03, 1.9348e-03, 1.9388e-03, 1.9379e-03, 1.9229e-03, 1.9323e-03],\n",
      "        [5.0680e-02, 5.0879e-02, 5.0821e-02, 5.0904e-02, 5.0654e-02, 5.0603e-02,\n",
      "         5.0566e-02, 5.0926e-02, 5.1024e-02, 5.0990e-02, 5.0592e-02, 5.0877e-02],\n",
      "        [2.5901e-01, 2.5989e-01, 2.5966e-01, 2.6007e-01, 2.5889e-01, 2.5864e-01,\n",
      "         2.5837e-01, 2.6015e-01, 2.6058e-01, 2.6052e-01, 2.5857e-01, 2.5989e-01],\n",
      "        [1.3496e-01, 1.3546e-01, 1.3532e-01, 1.3554e-01, 1.3489e-01, 1.3475e-01,\n",
      "         1.3464e-01, 1.3559e-01, 1.3584e-01, 1.3577e-01, 1.3472e-01, 1.3545e-01],\n",
      "        [5.8748e-03, 5.8969e-03, 5.8907e-03, 5.9002e-03, 5.8715e-03, 5.8657e-03,\n",
      "         5.8602e-03, 5.9026e-03, 5.9131e-03, 5.9100e-03, 5.8657e-03, 5.9005e-03],\n",
      "        [1.5410e-03, 1.5471e-03, 1.5453e-03, 1.5478e-03, 1.5401e-03, 1.5385e-03,\n",
      "         1.5376e-03, 1.5486e-03, 1.5516e-03, 1.5505e-03, 1.5383e-03, 1.5474e-03],\n",
      "        [1.2496e-02, 1.2545e-02, 1.2530e-02, 1.2551e-02, 1.2489e-02, 1.2477e-02,\n",
      "         1.2468e-02, 1.2556e-02, 1.2580e-02, 1.2572e-02, 1.2473e-02, 1.2542e-02],\n",
      "        [4.2455e-02, 4.2627e-02, 4.2578e-02, 4.2646e-02, 4.2433e-02, 4.2391e-02,\n",
      "         4.2369e-02, 4.2662e-02, 4.2748e-02, 4.2717e-02, 4.2385e-02, 4.2629e-02],\n",
      "        [2.1601e-02, 2.1686e-02, 2.1661e-02, 2.1696e-02, 2.1590e-02, 2.1568e-02,\n",
      "         2.1553e-02, 2.1706e-02, 2.1747e-02, 2.1733e-02, 2.1564e-02, 2.1685e-02],\n",
      "        [1.1408e-01, 1.1450e-01, 1.1438e-01, 1.1456e-01, 1.1402e-01, 1.1391e-01,\n",
      "         1.1381e-01, 1.1461e-01, 1.1481e-01, 1.1475e-01, 1.1388e-01, 1.1450e-01],\n",
      "        [1.4286e-03, 1.4352e-03, 1.4337e-03, 1.4363e-03, 1.4276e-03, 1.4255e-03,\n",
      "         1.4252e-03, 1.4379e-03, 1.4417e-03, 1.4398e-03, 1.4254e-03, 1.4352e-03],\n",
      "        [8.0823e-02, 8.1136e-02, 8.1047e-02, 8.1178e-02, 8.0784e-02, 8.0702e-02,\n",
      "         8.0637e-02, 8.1206e-02, 8.1364e-02, 8.1313e-02, 8.0680e-02, 8.1120e-02],\n",
      "        [4.6930e-02, 4.7115e-02, 4.7061e-02, 4.7137e-02, 4.6905e-02, 4.6857e-02,\n",
      "         4.6824e-02, 4.7159e-02, 4.7250e-02, 4.7218e-02, 4.6848e-02, 4.7115e-02],\n",
      "        [1.2191e-01, 1.2233e-01, 1.2222e-01, 1.2241e-01, 1.2185e-01, 1.2173e-01,\n",
      "         1.2161e-01, 1.2245e-01, 1.2266e-01, 1.2262e-01, 1.2170e-01, 1.2233e-01],\n",
      "        [9.2589e-02, 9.2891e-02, 9.2815e-02, 9.2959e-02, 9.2542e-02, 9.2435e-02,\n",
      "         9.2342e-02, 9.3000e-02, 9.3144e-02, 9.3137e-02, 9.2421e-02, 9.2884e-02],\n",
      "        [5.9620e-02, 5.9843e-02, 5.9781e-02, 5.9876e-02, 5.9592e-02, 5.9532e-02,\n",
      "         5.9479e-02, 5.9900e-02, 6.0009e-02, 5.9977e-02, 5.9517e-02, 5.9840e-02],\n",
      "        [6.7870e-02, 6.8079e-02, 6.8046e-02, 6.8147e-02, 6.7827e-02, 6.7745e-02,\n",
      "         6.7702e-02, 6.8179e-02, 6.8299e-02, 6.8292e-02, 6.7778e-02, 6.8146e-02],\n",
      "        [3.2973e-03, 3.3094e-03, 3.3062e-03, 3.3115e-03, 3.2961e-03, 3.2928e-03,\n",
      "         3.2897e-03, 3.3119e-03, 3.3183e-03, 3.3170e-03, 3.2914e-03, 3.3077e-03],\n",
      "        [9.5931e-02, 9.6307e-02, 9.6197e-02, 9.6355e-02, 9.5881e-02, 9.5783e-02,\n",
      "         9.5713e-02, 9.6394e-02, 9.6581e-02, 9.6517e-02, 9.5762e-02, 9.6298e-02],\n",
      "        [2.1555e-02, 2.1640e-02, 2.1615e-02, 2.1650e-02, 2.1544e-02, 2.1522e-02,\n",
      "         2.1506e-02, 2.1660e-02, 2.1701e-02, 2.1687e-02, 2.1517e-02, 2.1639e-02],\n",
      "        [1.6917e-02, 1.6983e-02, 1.6964e-02, 1.6991e-02, 1.6908e-02, 1.6891e-02,\n",
      "         1.6879e-02, 1.6998e-02, 1.7031e-02, 1.7020e-02, 1.6887e-02, 1.6981e-02],\n",
      "        [5.3445e-02, 5.3651e-02, 5.3592e-02, 5.3679e-02, 5.3420e-02, 5.3366e-02,\n",
      "         5.3323e-02, 5.3697e-02, 5.3800e-02, 5.3768e-02, 5.3351e-02, 5.3640e-02],\n",
      "        [1.3531e-02, 1.3584e-02, 1.3568e-02, 1.3590e-02, 1.3523e-02, 1.3509e-02,\n",
      "         1.3500e-02, 1.3597e-02, 1.3623e-02, 1.3614e-02, 1.3507e-02, 1.3585e-02],\n",
      "        [4.6216e-03, 4.6385e-03, 4.6341e-03, 4.6415e-03, 4.6199e-03, 4.6153e-03,\n",
      "         4.6109e-03, 4.6420e-03, 4.6509e-03, 4.6491e-03, 4.6133e-03, 4.6361e-03],\n",
      "        [9.4072e-03, 9.4373e-03, 9.4305e-03, 9.4445e-03, 9.4034e-03, 9.3941e-03,\n",
      "         9.3831e-03, 9.4468e-03, 9.4619e-03, 9.4624e-03, 9.3920e-03, 9.4371e-03],\n",
      "        [1.8459e-01, 1.8533e-01, 1.8512e-01, 1.8541e-01, 1.8450e-01, 1.8432e-01,\n",
      "         1.8421e-01, 1.8548e-01, 1.8585e-01, 1.8571e-01, 1.8428e-01, 1.8533e-01],\n",
      "        [4.9982e-03, 5.0171e-03, 5.0120e-03, 5.0200e-03, 4.9961e-03, 4.9911e-03,\n",
      "         4.9865e-03, 5.0210e-03, 5.0309e-03, 5.0283e-03, 4.9893e-03, 5.0151e-03],\n",
      "        [6.7328e-04, 6.7579e-04, 6.7528e-04, 6.7623e-04, 6.7285e-04, 6.7206e-04,\n",
      "         6.7169e-04, 6.7663e-04, 6.7791e-04, 6.7798e-04, 6.7219e-04, 6.7598e-04],\n",
      "        [3.9293e-01, 3.9427e-01, 3.9398e-01, 3.9456e-01, 3.9279e-01, 3.9246e-01,\n",
      "         3.9196e-01, 3.9458e-01, 3.9525e-01, 3.9527e-01, 3.9235e-01, 3.9432e-01],\n",
      "        [1.9566e-02, 1.9642e-02, 1.9620e-02, 1.9652e-02, 1.9557e-02, 1.9537e-02,\n",
      "         1.9522e-02, 1.9659e-02, 1.9697e-02, 1.9686e-02, 1.9532e-02, 1.9638e-02],\n",
      "        [4.1637e-01, 4.1808e-01, 4.1757e-01, 4.1826e-01, 4.1613e-01, 4.1571e-01,\n",
      "         4.1547e-01, 4.1844e-01, 4.1929e-01, 4.1899e-01, 4.1566e-01, 4.1809e-01],\n",
      "        [2.5098e-01, 2.5197e-01, 2.5167e-01, 2.5209e-01, 2.5084e-01, 2.5058e-01,\n",
      "         2.5041e-01, 2.5220e-01, 2.5269e-01, 2.5251e-01, 2.5054e-01, 2.5199e-01],\n",
      "        [1.4232e-01, 1.4289e-01, 1.4272e-01, 1.4296e-01, 1.4224e-01, 1.4210e-01,\n",
      "         1.4201e-01, 1.4302e-01, 1.4330e-01, 1.4321e-01, 1.4207e-01, 1.4288e-01],\n",
      "        [1.7956e-01, 1.8017e-01, 1.8001e-01, 1.8029e-01, 1.7947e-01, 1.7930e-01,\n",
      "         1.7912e-01, 1.8035e-01, 1.8065e-01, 1.8060e-01, 1.7925e-01, 1.8017e-01],\n",
      "        [3.4292e-02, 3.4428e-02, 3.4389e-02, 3.4445e-02, 3.4276e-02, 3.4241e-02,\n",
      "         3.4215e-02, 3.4459e-02, 3.4527e-02, 3.4505e-02, 3.4230e-02, 3.4418e-02],\n",
      "        [8.0648e-03, 8.0968e-03, 8.0872e-03, 8.1005e-03, 8.0601e-03, 8.0518e-03,\n",
      "         8.0466e-03, 8.1053e-03, 8.1207e-03, 8.1147e-03, 8.0507e-03, 8.0980e-03],\n",
      "        [2.2627e-03, 2.2711e-03, 2.2701e-03, 2.2737e-03, 2.2608e-03, 2.2591e-03,\n",
      "         2.2567e-03, 2.2745e-03, 2.2794e-03, 2.2810e-03, 2.2591e-03, 2.2732e-03],\n",
      "        [1.1984e-02, 1.2030e-02, 1.2017e-02, 1.2036e-02, 1.1979e-02, 1.1967e-02,\n",
      "         1.1957e-02, 1.2040e-02, 1.2063e-02, 1.2056e-02, 1.1963e-02, 1.2026e-02],\n",
      "        [8.5644e-02, 8.5948e-02, 8.5867e-02, 8.6001e-02, 8.5605e-02, 8.5519e-02,\n",
      "         8.5434e-02, 8.6032e-02, 8.6180e-02, 8.6149e-02, 8.5498e-02, 8.5944e-02],\n",
      "        [1.0534e-04, 1.0574e-04, 1.0567e-04, 1.0586e-04, 1.0524e-04, 1.0508e-04,\n",
      "         1.0508e-04, 1.0597e-04, 1.0620e-04, 1.0611e-04, 1.0517e-04, 1.0592e-04]],\n",
      "       device='cuda:0')}, 25: {'step': tensor(19527.), 'exp_avg': tensor([ 2.1296e-01,  3.4588e-02,  2.5311e-01,  1.8935e-01, -3.7162e-02,\n",
      "         8.2678e-02, -3.5175e-02,  2.5694e-01,  6.8344e-02, -5.3735e-02,\n",
      "         7.7940e-02,  1.2078e-01,  1.1763e-01,  8.5549e-02,  1.1380e-01,\n",
      "        -3.7351e-02,  2.3291e-01,  1.7263e-01, -2.6538e-02,  9.0531e-03,\n",
      "         3.8174e-02,  1.5949e-01, -4.8798e-03,  4.4833e-01, -1.1260e-02,\n",
      "         1.3611e-01,  2.1588e-01,  1.9738e-01, -3.6395e-02,  2.4938e-02,\n",
      "         6.6434e-02,  1.2778e-01,  8.7650e-02,  1.8006e-01,  3.0079e-02,\n",
      "        -1.5924e-01,  1.3351e-01,  1.5252e-01,  9.1252e-02,  1.3365e-01,\n",
      "         1.0507e-01, -2.3528e-02,  1.8310e-01,  9.0134e-02,  7.6334e-02,\n",
      "        -1.2874e-01,  7.3519e-02, -2.7054e-02,  2.7535e-02,  2.6728e-01,\n",
      "        -3.2971e-02,  7.8657e-03,  2.2748e-01, -7.6757e-02,  4.1193e-01,\n",
      "         3.0417e-01,  2.3733e-01,  1.8012e-01,  1.1248e-01,  5.9272e-02,\n",
      "        -2.6489e-04, -5.7747e-02,  1.4059e-01,  6.5091e-03], device='cuda:0'), 'exp_avg_sq': tensor([1.3513e+00, 2.4828e-02, 1.8286e+00, 1.1031e+00, 6.5698e-02, 1.8812e-01,\n",
      "        3.4958e-02, 2.0606e+00, 1.3246e-01, 7.1890e-02, 1.7049e-01, 4.0271e-01,\n",
      "        5.2420e-01, 1.8711e-01, 2.5362e-01, 9.0821e-02, 1.7216e+00, 8.0240e-01,\n",
      "        1.4731e-02, 4.2767e-03, 2.7462e-02, 6.9885e-01, 9.1635e-04, 6.9212e+00,\n",
      "        3.6019e-03, 5.0942e-01, 8.1289e-01, 9.1526e-01, 2.9586e-02, 1.9065e-02,\n",
      "        1.2953e-01, 5.6818e-01, 2.1672e-01, 7.3305e-01, 3.5720e-02, 6.6637e-01,\n",
      "        4.9780e-01, 3.7881e-01, 1.2361e-01, 3.9778e-01, 7.2171e-01, 1.3362e-02,\n",
      "        9.0405e-01, 2.1607e-01, 1.6133e-01, 4.3501e-01, 1.5007e-01, 1.7577e-02,\n",
      "        1.1027e-02, 2.1228e+00, 2.7405e-02, 5.3910e-03, 8.5013e-01, 1.5903e-01,\n",
      "        5.6579e+00, 2.6535e+00, 1.8063e+00, 5.8818e-01, 3.6052e-01, 9.8537e-02,\n",
      "        1.9410e-02, 8.5143e-02, 3.6430e-01, 2.3137e-03], device='cuda:0')}, 26: {'step': tensor(19527.), 'exp_avg': tensor([[[ 1.5211e-05,  2.0716e-06,  1.3434e-04,  ...,  6.6396e-06,\n",
      "           1.8374e-05,  3.6680e-05],\n",
      "         [-3.3191e-05,  3.1845e-04,  7.7388e-05,  ...,  8.2785e-06,\n",
      "           5.7002e-06, -2.6194e-05],\n",
      "         [ 1.7309e-05, -3.2980e-05,  1.4933e-04,  ...,  8.0302e-06,\n",
      "           1.3588e-06,  5.1244e-06],\n",
      "         ...,\n",
      "         [ 8.9135e-06,  8.1335e-06, -2.9231e-05,  ...,  8.7490e-05,\n",
      "           6.8790e-05,  1.0400e-04],\n",
      "         [-1.5888e-04, -1.1584e-04,  1.6540e-04,  ...,  5.2793e-07,\n",
      "          -1.4029e-05,  3.4818e-06],\n",
      "         [ 7.1070e-06, -3.0020e-04, -1.4417e-04,  ...,  1.1448e-04,\n",
      "          -1.4036e-05, -8.2809e-05]]], device='cuda:0'), 'exp_avg_sq': tensor([[[4.3470e-07, 2.4986e-09, 9.9331e-08,  ..., 2.0312e-07,\n",
      "          2.0702e-08, 4.4594e-09],\n",
      "         [2.7924e-07, 8.6786e-06, 4.2132e-08,  ..., 7.0325e-08,\n",
      "          2.4507e-07, 3.3955e-07],\n",
      "         [3.2767e-07, 1.6475e-07, 2.0406e-06,  ..., 2.0790e-07,\n",
      "          1.6875e-09, 2.2130e-08],\n",
      "         ...,\n",
      "         [2.1692e-08, 3.5839e-08, 5.9720e-08,  ..., 3.5819e-07,\n",
      "          4.3136e-08, 1.9000e-08],\n",
      "         [2.6680e-07, 5.5350e-07, 6.7109e-07,  ..., 1.7086e-09,\n",
      "          8.4903e-07, 3.2140e-08],\n",
      "         [1.1471e-08, 1.3350e-06, 4.1469e-07,  ..., 1.5711e-06,\n",
      "          2.5662e-07, 4.5480e-06]]], device='cuda:0')}, 27: {'step': tensor(19527.), 'exp_avg': tensor([[ 0.0114,  0.0101,  0.0110,  ...,  0.0127,  0.0043,  0.0076],\n",
      "        [ 0.0065,  0.0056,  0.0066,  ...,  0.0068,  0.0014,  0.0043],\n",
      "        [ 0.0073, -0.0009,  0.0021,  ...,  0.0062,  0.0070,  0.0076],\n",
      "        ...,\n",
      "        [-0.0016,  0.0003,  0.0008,  ..., -0.0018, -0.0042, -0.0014],\n",
      "        [ 0.0034,  0.0034,  0.0036,  ...,  0.0037,  0.0012,  0.0026],\n",
      "        [ 0.0029, -0.0002,  0.0003,  ...,  0.0015,  0.0022,  0.0017]],\n",
      "       device='cuda:0'), 'exp_avg_sq': tensor([[0.0014, 0.0008, 0.0012,  ..., 0.0014, 0.0006, 0.0005],\n",
      "        [0.0009, 0.0007, 0.0009,  ..., 0.0010, 0.0004, 0.0004],\n",
      "        [0.0019, 0.0011, 0.0016,  ..., 0.0020, 0.0007, 0.0009],\n",
      "        ...,\n",
      "        [0.0018, 0.0012, 0.0016,  ..., 0.0019, 0.0008, 0.0008],\n",
      "        [0.0004, 0.0003, 0.0005,  ..., 0.0005, 0.0003, 0.0002],\n",
      "        [0.0003, 0.0002, 0.0003,  ..., 0.0004, 0.0001, 0.0002]],\n",
      "       device='cuda:0')}, 28: {'step': tensor(19527.), 'exp_avg': tensor([[[-9.0410e-04,  1.7996e-04, -1.9389e-04,  ..., -1.5205e-04,\n",
      "           1.9362e-04, -4.9115e-05],\n",
      "         [-1.6061e-04,  6.0105e-05, -4.3854e-05,  ..., -3.2158e-05,\n",
      "           6.5553e-05, -1.4144e-05],\n",
      "         [-1.0536e-03,  1.1953e-04, -2.0099e-04,  ..., -1.5918e-04,\n",
      "           1.3305e-04, -4.6317e-05],\n",
      "         ...,\n",
      "         [ 2.1860e-04, -5.5700e-05,  5.0600e-05,  ...,  3.6868e-05,\n",
      "          -6.0610e-05,  1.1233e-05],\n",
      "         [-5.9611e-04,  2.5213e-04, -1.7000e-04,  ..., -1.2197e-04,\n",
      "           2.7512e-04, -5.4737e-05],\n",
      "         [-5.6791e-05,  8.6293e-06, -1.1101e-05,  ..., -1.0296e-05,\n",
      "           1.0604e-05, -4.1338e-06]],\n",
      "\n",
      "        [[ 9.8789e-06, -1.1521e-07,  2.2448e-06,  ...,  6.4372e-07,\n",
      "          -2.9860e-07,  2.6317e-07],\n",
      "         [ 1.0559e-06, -4.1424e-08,  2.5332e-07,  ...,  1.0969e-07,\n",
      "          -6.6133e-08,  4.5739e-08],\n",
      "         [ 1.2879e-05, -9.7411e-08,  2.8118e-06,  ...,  2.8448e-07,\n",
      "          -3.4911e-07,  2.9698e-07],\n",
      "         ...,\n",
      "         [-2.0295e-06,  5.0215e-08, -4.5495e-07,  ..., -1.7662e-07,\n",
      "           9.2005e-08, -3.5334e-08],\n",
      "         [ 3.8319e-06, -1.6030e-07,  9.7993e-07,  ...,  6.1003e-07,\n",
      "          -2.3509e-07,  1.6592e-07],\n",
      "         [ 5.1049e-07, -1.0674e-09,  1.0145e-07,  ...,  4.1097e-08,\n",
      "          -1.4777e-08,  1.8269e-08]],\n",
      "\n",
      "        [[ 8.6594e-04, -1.5816e-04,  1.8063e-04,  ...,  1.5090e-04,\n",
      "          -1.6768e-04,  4.7108e-05],\n",
      "         [ 1.5649e-04, -5.7365e-05,  4.2312e-05,  ...,  3.1971e-05,\n",
      "          -6.2924e-05,  1.3844e-05],\n",
      "         [ 1.0033e-03, -9.0736e-05,  1.8360e-04,  ...,  1.5763e-04,\n",
      "          -9.8238e-05,  4.3646e-05],\n",
      "         ...,\n",
      "         [-2.1021e-04,  5.0708e-05, -4.7625e-05,  ..., -3.6727e-05,\n",
      "           5.4899e-05, -1.0889e-05],\n",
      "         [ 5.8189e-04, -2.4370e-04,  1.6483e-04,  ...,  1.2173e-04,\n",
      "          -2.6552e-04,  5.3892e-05],\n",
      "         [ 5.4610e-05, -6.9935e-06,  1.0187e-05,  ...,  1.0077e-05,\n",
      "          -9.2041e-06,  3.9027e-06]]], device='cuda:0'), 'exp_avg_sq': tensor([[[1.1310e-04, 2.9161e-06, 6.1816e-06,  ..., 8.3589e-06,\n",
      "          7.5845e-06, 8.4425e-07],\n",
      "         [3.6786e-06, 2.0921e-07, 2.5637e-07,  ..., 2.5371e-07,\n",
      "          4.5830e-07, 3.1627e-08],\n",
      "         [1.3040e-04, 2.0443e-06, 6.3789e-06,  ..., 1.0293e-05,\n",
      "          6.1472e-06, 1.0012e-06],\n",
      "         ...,\n",
      "         [8.0511e-06, 2.7711e-07, 4.7030e-07,  ..., 5.6630e-07,\n",
      "          6.9985e-07, 5.3905e-08],\n",
      "         [5.2789e-05, 3.3622e-06, 3.7667e-06,  ..., 3.6644e-06,\n",
      "          7.4167e-06, 4.4788e-07],\n",
      "         [1.3499e-07, 2.1257e-09, 6.3203e-09,  ..., 1.1071e-08,\n",
      "          4.7247e-09, 1.5271e-09]],\n",
      "\n",
      "        [[1.9539e-09, 5.7955e-12, 6.2539e-11,  ..., 2.9800e-10,\n",
      "          1.3988e-11, 1.7920e-11],\n",
      "         [4.3418e-11, 2.2924e-13, 1.7924e-12,  ..., 5.1573e-12,\n",
      "          5.6886e-13, 3.7209e-13],\n",
      "         [2.6305e-09, 6.2989e-12, 7.8710e-11,  ..., 4.5247e-10,\n",
      "          1.4625e-11, 2.8571e-11],\n",
      "         ...,\n",
      "         [1.1442e-10, 4.3633e-13, 3.6064e-12,  ..., 1.5713e-11,\n",
      "          1.1321e-12, 7.2570e-13],\n",
      "         [5.3397e-10, 3.3283e-12, 2.1375e-11,  ..., 6.4313e-11,\n",
      "          8.7151e-12, 4.2760e-12],\n",
      "         [3.8747e-12, 6.6390e-15, 1.2987e-13,  ..., 7.8903e-13,\n",
      "          1.0182e-14, 8.2612e-14]],\n",
      "\n",
      "        [[1.0939e-04, 2.8190e-06, 6.0237e-06,  ..., 8.0506e-06,\n",
      "          7.3409e-06, 8.1517e-07],\n",
      "         [3.5863e-06, 2.0501e-07, 2.5144e-07,  ..., 2.4675e-07,\n",
      "          4.4916e-07, 3.0849e-08],\n",
      "         [1.2563e-04, 1.9518e-06, 6.1925e-06,  ..., 9.8692e-06,\n",
      "          5.8937e-06, 9.6124e-07],\n",
      "         ...,\n",
      "         [7.8169e-06, 2.6952e-07, 4.5992e-07,  ..., 5.4800e-07,\n",
      "          6.8113e-07, 5.2374e-08],\n",
      "         [5.1576e-05, 3.3012e-06, 3.7010e-06,  ..., 3.5718e-06,\n",
      "          7.2812e-06, 4.3792e-07],\n",
      "         [1.2905e-07, 2.0329e-09, 6.0889e-09,  ..., 1.0499e-08,\n",
      "          4.4988e-09, 1.4506e-09]]], device='cuda:0')}, 29: {'step': tensor(19527.), 'exp_avg': tensor([-4.9285e-01,  3.0600e-01, -1.5664e-01,  1.3818e-01,  3.6226e-01,\n",
      "        -1.5490e-01,  1.6892e-01,  6.2315e-01,  1.2932e-01,  8.3945e-02,\n",
      "        -1.1114e-01,  1.5893e-01,  1.1779e-01,  2.0294e-01, -1.5973e-01,\n",
      "         1.6123e-02,  5.2207e-02, -7.8535e-02,  1.8607e+00,  6.2391e-02,\n",
      "        -1.0586e+00, -1.0236e-01,  3.8248e-02,  7.9016e-02, -1.7844e-01,\n",
      "         1.7080e-02,  3.0115e-02, -6.9867e-03, -8.7094e-03,  3.4936e-02,\n",
      "         1.3167e+00,  7.7772e-02,  3.5675e-02, -6.0739e-02, -3.5379e-01,\n",
      "        -1.3265e-01, -4.9503e-01, -7.3054e-02,  7.1967e-02,  2.1847e-01,\n",
      "        -1.8801e-01,  1.4513e-01,  1.9079e-02,  5.5132e-02,  1.3827e-01,\n",
      "         1.4358e-01,  2.0629e-01, -2.3788e-01,  7.0499e-02, -2.0189e-01,\n",
      "        -8.0818e-02,  1.3527e-01,  1.9068e-01, -2.7025e-01,  1.1391e+00,\n",
      "         8.9526e-02, -1.4099e-01, -4.9673e-10, -9.9917e-02, -2.9221e-01,\n",
      "        -2.0660e-01, -1.1163e-01,  3.3211e-01, -3.9582e-02], device='cuda:0'), 'exp_avg_sq': tensor([3.0395e+01, 7.5191e+00, 3.2627e+00, 2.1989e+00, 1.4878e+01, 3.4988e+00,\n",
      "        2.2890e+00, 3.8758e+01, 2.0124e+00, 4.2891e-01, 1.5335e+00, 4.2802e+00,\n",
      "        1.3431e+00, 4.5748e+00, 2.7066e+00, 2.4863e-02, 1.4483e-01, 8.2066e-01,\n",
      "        1.9292e+02, 3.9661e-01, 1.4961e+02, 1.3835e+00, 6.5757e-02, 4.5516e-01,\n",
      "        4.0513e+00, 9.9252e-02, 1.5093e-01, 1.2900e-02, 5.7592e-03, 1.0046e-01,\n",
      "        1.7647e+02, 4.5050e-01, 1.4086e-01, 2.7896e-01, 1.7289e+01, 2.4202e+00,\n",
      "        2.9916e+01, 6.2194e-01, 9.7211e-01, 4.1338e+00, 4.6100e+00, 2.1141e+00,\n",
      "        6.5643e-02, 5.2918e-01, 1.7178e+00, 2.2143e+00, 4.3625e+00, 7.3499e+00,\n",
      "        6.0228e-01, 5.1966e+00, 7.3447e-01, 1.9039e+00, 3.8001e+00, 9.4276e+00,\n",
      "        1.3828e+02, 7.9783e-01, 2.0466e+00, 3.6607e-13, 1.2378e+00, 1.0914e+01,\n",
      "        6.1198e+00, 1.7581e+00, 1.2523e+01, 2.6471e-01], device='cuda:0')}, 30: {'step': tensor(19527.), 'exp_avg': tensor([[[[ 4.6811e-05,  1.5218e-04, -4.0322e-05]],\n",
      "\n",
      "         [[ 7.4493e-05,  1.4841e-04, -2.4246e-04]],\n",
      "\n",
      "         [[ 7.9904e-05,  2.2500e-04, -1.1213e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 2.3492e-04,  6.1847e-04, -4.0545e-04]],\n",
      "\n",
      "         [[ 6.8389e-05,  1.1701e-04, -1.6183e-04]],\n",
      "\n",
      "         [[ 2.2935e-04,  5.6035e-04, -4.0461e-04]]],\n",
      "\n",
      "\n",
      "        [[[-7.9213e-03, -7.8075e-03, -6.0856e-03]],\n",
      "\n",
      "         [[-4.0688e-03, -3.8778e-03, -3.3689e-03]],\n",
      "\n",
      "         [[-1.0696e-02, -1.0651e-02, -8.5066e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-2.5782e-02, -2.5884e-02, -2.1135e-02]],\n",
      "\n",
      "         [[-3.2664e-03, -3.1363e-03, -2.6940e-03]],\n",
      "\n",
      "         [[-2.2954e-02, -2.3050e-02, -1.8948e-02]]],\n",
      "\n",
      "\n",
      "        [[[ 9.5304e-03,  1.0126e-02,  8.8892e-03]],\n",
      "\n",
      "         [[ 4.5914e-03,  4.7648e-03,  4.7248e-03]],\n",
      "\n",
      "         [[ 1.2053e-02,  1.2646e-02,  1.1386e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 2.7512e-02,  2.8445e-02,  2.6294e-02]],\n",
      "\n",
      "         [[ 3.6089e-03,  3.7128e-03,  3.6366e-03]],\n",
      "\n",
      "         [[ 2.4187e-02,  2.4978e-02,  2.3343e-02]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 6.0531e-03,  5.3183e-03,  3.8797e-03]],\n",
      "\n",
      "         [[-3.1097e-03, -3.0933e-03, -2.7539e-03]],\n",
      "\n",
      "         [[ 4.6011e-03,  4.0038e-03,  2.6503e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 2.1234e-03,  1.4651e-03, -2.7138e-04]],\n",
      "\n",
      "         [[-2.3806e-03, -2.3551e-03, -2.1048e-03]],\n",
      "\n",
      "         [[-3.5165e-04, -6.9112e-04, -1.9253e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 7.9275e-04,  8.5806e-04,  7.0848e-04]],\n",
      "\n",
      "         [[-3.4540e-04, -3.8530e-04, -4.2869e-04]],\n",
      "\n",
      "         [[ 8.0679e-04,  8.2460e-04,  6.4236e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.3696e-03,  1.2596e-03,  8.5344e-04]],\n",
      "\n",
      "         [[-2.2267e-04, -2.5687e-04, -2.9433e-04]],\n",
      "\n",
      "         [[ 1.1284e-03,  9.9658e-04,  6.5077e-04]]],\n",
      "\n",
      "\n",
      "        [[[-2.3690e-04, -1.6356e-05,  6.2105e-04]],\n",
      "\n",
      "         [[ 6.7634e-05,  1.6991e-04, -1.3932e-05]],\n",
      "\n",
      "         [[-2.7313e-04, -2.0564e-05,  4.9962e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-5.7756e-04, -5.4517e-05,  3.5404e-04]],\n",
      "\n",
      "         [[ 5.0220e-05,  1.2142e-04, -8.2780e-06]],\n",
      "\n",
      "         [[-4.5645e-04,  5.4068e-06,  1.5193e-04]]]], device='cuda:0'), 'exp_avg_sq': tensor([[[[4.6814e-05, 3.6916e-05, 1.7080e-05]],\n",
      "\n",
      "         [[9.2376e-05, 1.7278e-04, 4.2179e-05]],\n",
      "\n",
      "         [[1.1912e-04, 1.2025e-04, 4.6239e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[9.0352e-04, 1.0659e-03, 3.6550e-04]],\n",
      "\n",
      "         [[4.9686e-05, 8.5680e-05, 2.2382e-05]],\n",
      "\n",
      "         [[8.1111e-04, 9.9484e-04, 3.3231e-04]]],\n",
      "\n",
      "\n",
      "        [[[1.6727e-02, 1.7521e-02, 1.1998e-02]],\n",
      "\n",
      "         [[3.3930e-03, 3.8713e-03, 2.7152e-03]],\n",
      "\n",
      "         [[2.5238e-02, 2.6950e-02, 1.8761e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[1.1843e-01, 1.2911e-01, 9.1096e-02]],\n",
      "\n",
      "         [[2.1952e-03, 2.4805e-03, 1.7575e-03]],\n",
      "\n",
      "         [[9.3307e-02, 1.0242e-01, 7.2679e-02]]],\n",
      "\n",
      "\n",
      "        [[[1.0502e-02, 1.3130e-02, 1.2380e-02]],\n",
      "\n",
      "         [[1.6316e-03, 1.9630e-03, 1.7127e-03]],\n",
      "\n",
      "         [[1.6253e-02, 1.9889e-02, 1.8718e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[7.7847e-02, 9.2889e-02, 8.6392e-02]],\n",
      "\n",
      "         [[1.1079e-03, 1.3166e-03, 1.1689e-03]],\n",
      "\n",
      "         [[6.1144e-02, 7.2565e-02, 6.7489e-02]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[5.4762e-03, 4.8565e-03, 3.1320e-03]],\n",
      "\n",
      "         [[1.4230e-03, 1.6151e-03, 1.3525e-03]],\n",
      "\n",
      "         [[4.7856e-03, 4.5045e-03, 3.0823e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[1.3802e-02, 1.5378e-02, 1.2553e-02]],\n",
      "\n",
      "         [[8.1338e-04, 9.3046e-04, 7.8681e-04]],\n",
      "\n",
      "         [[1.0840e-02, 1.2704e-02, 1.0842e-02]]],\n",
      "\n",
      "\n",
      "        [[[5.1135e-04, 5.2066e-04, 3.6816e-04]],\n",
      "\n",
      "         [[1.4663e-05, 1.6236e-05, 1.0390e-05]],\n",
      "\n",
      "         [[6.4521e-04, 6.5328e-04, 4.6097e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[2.3183e-03, 2.3141e-03, 1.6220e-03]],\n",
      "\n",
      "         [[1.2018e-05, 1.2440e-05, 8.1454e-06]],\n",
      "\n",
      "         [[1.6559e-03, 1.6526e-03, 1.1573e-03]]],\n",
      "\n",
      "\n",
      "        [[[3.1406e-05, 6.9218e-06, 3.0760e-05]],\n",
      "\n",
      "         [[9.9973e-06, 3.1153e-05, 1.1947e-06]],\n",
      "\n",
      "         [[4.4140e-05, 2.1677e-05, 3.6236e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[2.0495e-04, 1.9037e-04, 1.2304e-04]],\n",
      "\n",
      "         [[5.7353e-06, 1.5386e-05, 7.9449e-07]],\n",
      "\n",
      "         [[1.6307e-04, 1.7692e-04, 8.6501e-05]]]], device='cuda:0')}, 31: {'step': tensor(19527.), 'exp_avg': tensor([ 1.1709e-03, -4.8357e-02,  5.2793e-02,  5.8903e-03, -1.9039e-02,\n",
      "         1.6658e-03,  6.9538e-03,  1.2017e-07,  1.2386e-07, -3.8100e-02,\n",
      "         7.8833e-03,  2.0043e-02,  2.5089e-02,  2.4982e-02,  9.6889e-03,\n",
      "        -3.0399e-02, -4.8617e-02,  2.9422e-02,  2.6812e-04, -2.0853e-02,\n",
      "         2.6127e-03,  1.7086e-03,  1.0129e-02,  1.8192e-02, -9.7686e-02,\n",
      "         1.8384e-02, -3.9017e-02,  1.8672e-02,  6.8456e-03, -2.8423e-02,\n",
      "         2.8848e-02,  2.3499e-02, -2.7502e-08, -1.8226e-02, -2.3154e-03,\n",
      "         2.6286e-02,  4.0073e-02, -1.6038e-02, -1.5570e-02,  1.4027e-02,\n",
      "         1.3521e-02,  2.2383e-02,  5.8448e-02,  5.8396e-04,  6.0365e-07,\n",
      "        -1.6431e-02,  2.6661e-02,  1.8085e-02,  1.0104e-02,  3.3626e-02,\n",
      "         2.5183e-02, -5.1663e-02, -3.3535e-02,  7.4684e-03, -5.5596e-02,\n",
      "         1.9512e-02,  1.2086e-02, -1.4922e-02,  1.1400e-02,  3.8906e-02,\n",
      "         3.7087e-03,  1.6438e-04,  2.2325e-03, -6.5400e-05], device='cuda:0'), 'exp_avg_sq': tensor([3.8897e-03, 4.3555e-01, 3.1010e-01, 1.2035e-02, 1.0816e-01, 3.9597e-03,\n",
      "        1.4057e-02, 1.1205e-09, 4.1350e-09, 2.8843e-01, 2.1153e-02, 6.1694e-02,\n",
      "        9.3388e-02, 5.7705e-02, 5.2237e-03, 1.5878e-01, 4.3708e-01, 1.1392e-01,\n",
      "        1.1753e-02, 8.3196e-02, 3.3912e-02, 3.6381e-04, 5.1390e-02, 5.9088e-02,\n",
      "        1.1465e+00, 5.7053e-02, 3.1072e-01, 4.6056e-02, 1.8824e-02, 1.2401e-01,\n",
      "        1.3202e-01, 1.0111e-01, 3.8295e-07, 1.3814e-02, 1.0266e-03, 1.7438e-02,\n",
      "        2.1195e-01, 3.3137e-02, 1.0217e-01, 3.6334e-02, 3.5982e-02, 1.3661e-01,\n",
      "        3.7957e-01, 4.3470e-04, 6.0966e-09, 7.6973e-02, 1.6648e-01, 6.5184e-02,\n",
      "        3.8372e-02, 1.1624e-01, 5.5079e-02, 4.2907e-01, 3.0424e-01, 8.5340e-03,\n",
      "        5.1339e-01, 1.2670e-02, 3.4179e-02, 5.1998e-02, 2.6392e-02, 9.3834e-02,\n",
      "        1.2887e-03, 5.1976e-02, 7.3720e-03, 6.9550e-04], device='cuda:0')}, 32: {'step': tensor(19527.), 'exp_avg': tensor([[[[ 1.3014e-03]],\n",
      "\n",
      "         [[ 6.3431e-04]],\n",
      "\n",
      "         [[-2.1937e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-6.8336e-04]],\n",
      "\n",
      "         [[ 2.9850e-03]],\n",
      "\n",
      "         [[ 3.2600e-05]]],\n",
      "\n",
      "\n",
      "        [[[-5.3748e-02]],\n",
      "\n",
      "         [[-1.3082e-02]],\n",
      "\n",
      "         [[-5.1300e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.5965e-02]],\n",
      "\n",
      "         [[-5.4837e-02]],\n",
      "\n",
      "         [[-2.5727e-03]]],\n",
      "\n",
      "\n",
      "        [[[ 5.2651e-02]],\n",
      "\n",
      "         [[ 1.5100e-02]],\n",
      "\n",
      "         [[ 4.5166e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-1.4911e-02]],\n",
      "\n",
      "         [[ 6.0656e-02]],\n",
      "\n",
      "         [[ 3.5305e-03]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 6.2644e-03]],\n",
      "\n",
      "         [[-4.4913e-03]],\n",
      "\n",
      "         [[ 2.2444e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 6.5463e-04]],\n",
      "\n",
      "         [[-2.0461e-02]],\n",
      "\n",
      "         [[ 5.1510e-04]]],\n",
      "\n",
      "\n",
      "        [[[ 3.8678e-03]],\n",
      "\n",
      "         [[-6.6372e-05]],\n",
      "\n",
      "         [[ 6.2551e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-8.7416e-04]],\n",
      "\n",
      "         [[-4.4003e-04]],\n",
      "\n",
      "         [[ 5.6137e-05]]],\n",
      "\n",
      "\n",
      "        [[[-1.8642e-05]],\n",
      "\n",
      "         [[-8.9854e-05]],\n",
      "\n",
      "         [[ 2.3164e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-1.1236e-04]],\n",
      "\n",
      "         [[-9.5642e-05]],\n",
      "\n",
      "         [[-8.4178e-05]]]], device='cuda:0'), 'exp_avg_sq': tensor([[[[1.1289e-03]],\n",
      "\n",
      "         [[2.4085e-04]],\n",
      "\n",
      "         [[1.9495e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[1.5009e-04]],\n",
      "\n",
      "         [[4.5069e-03]],\n",
      "\n",
      "         [[8.9558e-06]]],\n",
      "\n",
      "\n",
      "        [[[5.2670e-01]],\n",
      "\n",
      "         [[2.8701e-02]],\n",
      "\n",
      "         [[4.4757e-01]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.2001e-02]],\n",
      "\n",
      "         [[5.0056e-01]],\n",
      "\n",
      "         [[1.7790e-04]]],\n",
      "\n",
      "\n",
      "        [[[4.0538e-01]],\n",
      "\n",
      "         [[1.9369e-02]],\n",
      "\n",
      "         [[3.8387e-01]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[3.3061e-02]],\n",
      "\n",
      "         [[2.9865e-01]],\n",
      "\n",
      "         [[4.0185e-04]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[5.7360e-02]],\n",
      "\n",
      "         [[5.7531e-03]],\n",
      "\n",
      "         [[4.9860e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.5368e-03]],\n",
      "\n",
      "         [[9.7354e-02]],\n",
      "\n",
      "         [[1.0088e-04]]],\n",
      "\n",
      "\n",
      "        [[[1.0415e-02]],\n",
      "\n",
      "         [[3.0777e-04]],\n",
      "\n",
      "         [[1.2511e-02]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[8.9092e-04]],\n",
      "\n",
      "         [[5.2190e-03]],\n",
      "\n",
      "         [[4.3063e-06]]],\n",
      "\n",
      "\n",
      "        [[[1.5846e-04]],\n",
      "\n",
      "         [[2.0303e-05]],\n",
      "\n",
      "         [[2.2925e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[3.8018e-05]],\n",
      "\n",
      "         [[5.4878e-04]],\n",
      "\n",
      "         [[7.0045e-07]]]], device='cuda:0')}, 33: {'step': tensor(19527.), 'exp_avg': tensor([ 1.1709e-03, -4.8357e-02,  5.2792e-02,  5.8901e-03, -1.9039e-02,\n",
      "         1.6660e-03,  6.9537e-03, -3.2148e-08,  2.5542e-08, -3.8100e-02,\n",
      "         7.8830e-03,  2.0043e-02,  2.5090e-02,  2.4982e-02,  9.6889e-03,\n",
      "        -3.0400e-02, -4.8616e-02,  2.9422e-02,  2.6790e-04, -2.0853e-02,\n",
      "         2.6126e-03,  1.7078e-03,  1.0128e-02,  1.8192e-02, -9.7686e-02,\n",
      "         1.8383e-02, -3.9018e-02,  1.8673e-02,  6.8461e-03, -2.8423e-02,\n",
      "         2.8847e-02,  2.3499e-02, -2.7408e-08, -1.8225e-02, -2.3152e-03,\n",
      "         2.6286e-02,  4.0074e-02, -1.6037e-02, -1.5570e-02,  1.4027e-02,\n",
      "         1.3521e-02,  2.2382e-02,  5.8448e-02,  5.8402e-04,  7.0812e-08,\n",
      "        -1.6431e-02,  2.6661e-02,  1.8085e-02,  1.0105e-02,  3.3626e-02,\n",
      "         2.5183e-02, -5.1663e-02, -3.3536e-02,  7.4689e-03, -5.5596e-02,\n",
      "         1.9513e-02,  1.2085e-02, -1.4922e-02,  1.1400e-02,  3.8906e-02,\n",
      "         3.7087e-03,  1.6449e-04,  2.2330e-03, -6.5498e-05], device='cuda:0'), 'exp_avg_sq': tensor([3.8897e-03, 4.3555e-01, 3.1010e-01, 1.2035e-02, 1.0816e-01, 3.9597e-03,\n",
      "        1.4057e-02, 1.1205e-09, 4.1350e-09, 2.8843e-01, 2.1153e-02, 6.1694e-02,\n",
      "        9.3388e-02, 5.7705e-02, 5.2237e-03, 1.5878e-01, 4.3708e-01, 1.1392e-01,\n",
      "        1.1753e-02, 8.3196e-02, 3.3912e-02, 3.6381e-04, 5.1390e-02, 5.9088e-02,\n",
      "        1.1465e+00, 5.7053e-02, 3.1072e-01, 4.6056e-02, 1.8824e-02, 1.2401e-01,\n",
      "        1.3202e-01, 1.0111e-01, 3.8295e-07, 1.3814e-02, 1.0266e-03, 1.7438e-02,\n",
      "        2.1195e-01, 3.3137e-02, 1.0217e-01, 3.6334e-02, 3.5982e-02, 1.3661e-01,\n",
      "        3.7957e-01, 4.3470e-04, 6.0965e-09, 7.6973e-02, 1.6648e-01, 6.5184e-02,\n",
      "        3.8372e-02, 1.1624e-01, 5.5079e-02, 4.2907e-01, 3.0424e-01, 8.5340e-03,\n",
      "        5.1339e-01, 1.2670e-02, 3.4179e-02, 5.1998e-02, 2.6392e-02, 9.3834e-02,\n",
      "        1.2887e-03, 5.1976e-02, 7.3720e-03, 6.9550e-04], device='cuda:0')}, 34: {'step': tensor(19527.), 'exp_avg': tensor([-9.5257e-02,  1.1679e-01,  1.7487e-02, -3.9166e-02,  6.3162e-02,\n",
      "         7.0202e-02, -7.5000e-02,  5.1927e-02,  6.6982e-02,  1.0720e-01,\n",
      "         7.5224e-03,  1.8101e-02,  8.1899e-02, -4.5254e-02, -8.5763e-02,\n",
      "         2.1133e-02, -1.7096e-02,  6.1487e-02,  6.2595e-02,  1.8297e-02,\n",
      "        -1.7237e-02, -6.3191e-02, -5.7179e-02, -4.6391e-02, -5.4061e-02,\n",
      "        -5.9342e-02, -1.0200e-01,  4.0937e-02, -2.1369e-02,  1.0932e-02,\n",
      "         3.3641e-02,  2.9530e-02,  5.6132e-02, -3.8721e-02, -6.7254e-02,\n",
      "        -1.0471e-01,  3.9354e-02, -6.8407e-02,  1.3534e-01,  2.0391e-02,\n",
      "        -2.9017e-02,  3.3150e-02,  4.7082e-02, -5.7835e-02,  5.7562e-02,\n",
      "         5.7584e-02,  8.8145e-03,  5.9162e-02, -4.4283e-02, -3.0084e-02,\n",
      "        -5.4076e-02, -5.0154e-02, -6.6998e-02, -4.9359e-02, -8.1792e-02,\n",
      "        -9.7128e-02, -2.5566e-02,  1.6769e-02, -1.3001e-02, -6.4749e-02,\n",
      "         6.8922e-02,  1.1410e-01,  1.1444e-04,  4.8517e-02], device='cuda:0'), 'exp_avg_sq': tensor([1.4318, 1.3953, 0.1141, 0.2448, 0.4254, 0.5905, 0.5934, 0.3302, 0.5575,\n",
      "        1.7293, 0.0292, 0.0833, 1.1883, 0.2599, 0.8627, 0.0400, 0.1679, 0.6224,\n",
      "        0.5403, 0.0269, 0.0748, 0.5673, 0.5981, 0.3719, 0.4590, 0.5172, 1.5053,\n",
      "        0.2762, 0.0277, 0.0313, 0.2653, 0.3698, 0.3846, 0.2233, 0.4784, 1.0828,\n",
      "        0.2710, 0.2990, 3.4634, 0.0977, 0.1198, 0.6744, 0.3561, 0.4833, 0.3954,\n",
      "        0.5303, 0.2467, 0.7275, 0.3030, 0.1432, 0.4133, 0.2425, 0.2279, 0.3706,\n",
      "        0.7518, 0.8659, 0.0851, 0.0218, 0.0273, 0.6364, 0.4580, 2.3925, 0.0236,\n",
      "        0.3203], device='cuda:0')}, 35: {'step': tensor(19527.), 'exp_avg': tensor([ 0.1164,  0.1120,  0.0799,  0.0588,  0.0584, -0.0498,  0.0936, -0.0585,\n",
      "        -0.0755,  0.1025,  0.0455,  0.0544,  0.0540,  0.0760,  0.0985,  0.0660,\n",
      "        -0.0888,  0.0528, -0.0416,  0.0532,  0.0482,  0.0733,  0.0925,  0.0853,\n",
      "        -0.1264,  0.0915, -0.1007,  0.0634,  0.0519,  0.0609,  0.0636,  0.0549,\n",
      "        -0.0633,  0.0530,  0.0741,  0.1206,  0.0660,  0.0660,  0.0775,  0.0536,\n",
      "         0.0554,  0.0631,  0.0715,  0.0677, -0.0649,  0.0551,  0.0731,  0.0471,\n",
      "         0.0744,  0.0901,  0.0853, -0.1007, -0.0662,  0.0667, -0.0991,  0.1142,\n",
      "         0.0525,  0.0565,  0.0498,  0.1085, -0.0628,  0.0494,  0.0501, -0.0547],\n",
      "       device='cuda:0'), 'exp_avg_sq': tensor([1.9191, 1.5103, 0.7185, 0.4531, 0.4032, 0.2476, 0.8444, 0.4059, 0.6849,\n",
      "        1.2924, 0.2643, 0.3487, 0.3708, 0.5766, 1.1075, 0.4992, 0.9143, 0.3262,\n",
      "        0.1913, 0.3283, 0.3802, 0.7175, 1.0926, 0.8842, 1.7473, 0.9622, 1.2779,\n",
      "        0.4672, 0.3328, 0.4226, 0.4867, 0.3470, 0.4728, 0.3288, 0.6603, 1.4235,\n",
      "        0.5080, 0.7305, 0.7805, 0.3501, 0.3843, 0.5252, 0.6267, 0.6187, 0.4862,\n",
      "        0.3858, 0.6764, 0.2712, 0.6909, 0.8144, 0.7509, 1.1923, 0.5331, 0.5696,\n",
      "        1.1833, 1.1356, 0.3199, 0.3812, 0.2951, 1.1587, 0.3878, 0.3402, 0.3121,\n",
      "        0.3679], device='cuda:0')}, 36: {'step': tensor(19527.), 'exp_avg': tensor([[[[ 3.0743e-05, -1.2932e-03, -9.9145e-04,  ..., -4.6443e-04,\n",
      "           -6.1288e-04,  5.1409e-04]],\n",
      "\n",
      "         [[ 9.0853e-05, -9.9659e-04, -1.0839e-03,  ..., -1.3560e-04,\n",
      "           -7.3319e-04,  2.8933e-04]],\n",
      "\n",
      "         [[ 2.5046e-05, -1.2133e-03, -1.4247e-04,  ...,  1.6625e-04,\n",
      "            1.9130e-06, -1.9435e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-2.7917e-05, -1.1884e-03, -3.5392e-04,  ..., -1.3198e-04,\n",
      "           -4.4184e-04,  1.2162e-04]],\n",
      "\n",
      "         [[-4.7483e-05, -1.2846e-03, -3.7307e-04,  ..., -1.1884e-04,\n",
      "           -6.1216e-04,  1.8679e-04]],\n",
      "\n",
      "         [[-3.4402e-04, -1.0646e-03,  3.8397e-04,  ...,  6.5989e-07,\n",
      "           -3.7766e-05, -8.4465e-04]]],\n",
      "\n",
      "\n",
      "        [[[ 2.3225e-05, -6.9845e-04, -9.9746e-04,  ..., -3.3828e-04,\n",
      "           -3.9748e-04,  4.0836e-04]],\n",
      "\n",
      "         [[ 1.7698e-05, -7.9208e-04, -8.5585e-04,  ..., -1.0795e-04,\n",
      "           -4.7577e-04,  2.5182e-04]],\n",
      "\n",
      "         [[ 5.3175e-05, -9.7600e-04, -6.8653e-05,  ...,  1.3487e-04,\n",
      "            1.5222e-04, -1.7767e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 2.8491e-05, -9.6721e-04, -8.3660e-05,  ..., -3.2801e-05,\n",
      "           -8.1577e-05, -3.3508e-05]],\n",
      "\n",
      "         [[ 1.7525e-05, -8.6360e-04, -4.4969e-04,  ..., -1.1340e-04,\n",
      "           -4.4898e-04,  1.9358e-04]],\n",
      "\n",
      "         [[-1.6911e-04, -7.8589e-04,  2.2341e-04,  ...,  3.8338e-05,\n",
      "           -1.2619e-05, -5.8395e-04]]],\n",
      "\n",
      "\n",
      "        [[[ 1.5791e-04, -6.0489e-04,  5.9340e-04,  ...,  6.7700e-04,\n",
      "            1.4694e-03, -8.0416e-04]],\n",
      "\n",
      "         [[ 1.6755e-04, -6.9686e-04,  4.0052e-04,  ...,  2.7948e-04,\n",
      "            1.4003e-03, -7.4996e-04]],\n",
      "\n",
      "         [[ 1.2363e-04, -9.2700e-04,  1.1861e-03,  ...,  5.1536e-04,\n",
      "            2.0242e-03, -1.1766e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 9.7604e-05, -9.6845e-04,  1.2116e-03,  ...,  3.7695e-04,\n",
      "            1.6910e-03, -1.0144e-03]],\n",
      "\n",
      "         [[ 3.7978e-05, -7.5260e-04,  6.1467e-04,  ...,  3.6902e-04,\n",
      "            1.2812e-03, -7.7295e-04]],\n",
      "\n",
      "         [[-1.2880e-04, -3.0107e-04,  1.2465e-03,  ...,  3.5276e-04,\n",
      "            2.0688e-03, -1.2013e-03]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[ 2.0442e-04,  4.4563e-04,  1.4012e-03,  ...,  1.0986e-03,\n",
      "            1.9269e-03, -1.1990e-03]],\n",
      "\n",
      "         [[ 1.1979e-04,  2.3483e-04,  1.3045e-03,  ...,  5.7051e-04,\n",
      "            2.0413e-03, -1.1302e-03]],\n",
      "\n",
      "         [[ 9.6672e-05,  1.1458e-04,  2.1916e-03,  ...,  8.4949e-04,\n",
      "            2.7654e-03, -1.5951e-03]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 1.3050e-05, -7.7778e-05,  1.9285e-03,  ...,  6.7616e-04,\n",
      "            2.4456e-03, -1.3815e-03]],\n",
      "\n",
      "         [[ 6.7718e-05,  4.2768e-05,  2.1016e-03,  ...,  6.7459e-04,\n",
      "            2.4537e-03, -1.3747e-03]],\n",
      "\n",
      "         [[ 2.9517e-04,  3.5694e-04,  1.7125e-03,  ...,  8.2481e-04,\n",
      "            2.6157e-03, -3.6517e-04]]],\n",
      "\n",
      "\n",
      "        [[[ 4.7166e-04,  1.0601e-03, -9.1886e-05,  ...,  4.6165e-04,\n",
      "            4.5634e-04, -2.7747e-04]],\n",
      "\n",
      "         [[ 1.8052e-05,  9.1659e-04,  1.8064e-04,  ...,  3.5554e-04,\n",
      "            6.9536e-04, -3.4651e-04]],\n",
      "\n",
      "         [[ 5.9457e-06,  6.6402e-04,  1.3948e-03,  ...,  6.7485e-04,\n",
      "            1.4383e-03, -8.4315e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 7.1213e-06,  7.0263e-04,  1.0145e-03,  ...,  5.1324e-04,\n",
      "            1.2681e-03, -6.7846e-04]],\n",
      "\n",
      "         [[ 1.8900e-05,  6.6419e-04,  1.1490e-03,  ...,  4.8074e-04,\n",
      "            1.2175e-03, -6.5248e-04]],\n",
      "\n",
      "         [[ 4.4900e-04,  9.0389e-04,  8.1279e-04,  ...,  7.2156e-04,\n",
      "            1.3138e-03,  6.6090e-04]]],\n",
      "\n",
      "\n",
      "        [[[ 4.0402e-04,  6.1803e-04, -9.3611e-04,  ...,  4.8774e-05,\n",
      "           -2.3629e-04,  2.4124e-04]],\n",
      "\n",
      "         [[ 1.6346e-04,  3.4294e-04, -6.4320e-04,  ...,  2.2239e-04,\n",
      "           -2.0877e-04,  5.2750e-05]],\n",
      "\n",
      "         [[ 9.3097e-05,  2.3525e-04,  9.9433e-04,  ...,  6.4296e-04,\n",
      "            1.0248e-03, -6.7574e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[-3.8964e-05,  2.8241e-04, -3.3312e-05,  ...,  1.7245e-04,\n",
      "            3.8207e-04, -1.7724e-04]],\n",
      "\n",
      "         [[-1.6855e-05, -1.7623e-05,  6.0367e-04,  ...,  8.7752e-05,\n",
      "            4.6932e-04, -2.5416e-04]],\n",
      "\n",
      "         [[ 2.5457e-04,  1.6055e-04,  3.4120e-04,  ...,  4.5464e-04,\n",
      "            3.6906e-04,  6.3041e-04]]]], device='cuda:0'), 'exp_avg_sq': tensor([[[[1.4008e-06, 3.0828e-05, 2.9049e-04,  ..., 7.7260e-05,\n",
      "           3.0302e-04, 1.0929e-04]],\n",
      "\n",
      "         [[5.7019e-07, 2.5037e-05, 1.8765e-04,  ..., 1.2540e-05,\n",
      "           2.4003e-04, 7.7008e-05]],\n",
      "\n",
      "         [[5.4667e-07, 2.4732e-05, 1.9149e-04,  ..., 1.2697e-05,\n",
      "           2.4375e-04, 7.8106e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.0733e-07, 2.2664e-05, 1.7594e-04,  ..., 1.1907e-05,\n",
      "           2.4710e-04, 7.7348e-05]],\n",
      "\n",
      "         [[5.0156e-07, 2.2927e-05, 1.8294e-04,  ..., 1.1864e-05,\n",
      "           2.4344e-04, 7.6473e-05]],\n",
      "\n",
      "         [[2.7827e-06, 1.5538e-05, 9.1817e-05,  ..., 9.4680e-06,\n",
      "           2.5775e-04, 4.8519e-05]]],\n",
      "\n",
      "\n",
      "        [[[1.5319e-06, 3.0780e-05, 3.1355e-04,  ..., 8.0610e-05,\n",
      "           3.1335e-04, 1.1336e-04]],\n",
      "\n",
      "         [[6.2615e-07, 2.5674e-05, 2.0657e-04,  ..., 1.3112e-05,\n",
      "           2.4961e-04, 8.0332e-05]],\n",
      "\n",
      "         [[6.2618e-07, 2.5396e-05, 2.0815e-04,  ..., 1.3343e-05,\n",
      "           2.5223e-04, 8.1129e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.3770e-07, 2.2737e-05, 1.9243e-04,  ..., 1.2422e-05,\n",
      "           2.5748e-04, 8.0868e-05]],\n",
      "\n",
      "         [[5.1517e-07, 2.3200e-05, 1.9863e-04,  ..., 1.2168e-05,\n",
      "           2.5414e-04, 7.9773e-05]],\n",
      "\n",
      "         [[2.4434e-06, 1.6450e-05, 9.9380e-05,  ..., 9.3007e-06,\n",
      "           2.6970e-04, 4.9825e-05]]],\n",
      "\n",
      "\n",
      "        [[[1.5854e-06, 3.1032e-05, 3.3289e-04,  ..., 8.3995e-05,\n",
      "           3.2335e-04, 1.1733e-04]],\n",
      "\n",
      "         [[6.4668e-07, 2.6031e-05, 2.2296e-04,  ..., 1.3825e-05,\n",
      "           2.5949e-04, 8.3655e-05]],\n",
      "\n",
      "         [[6.3463e-07, 2.5985e-05, 2.2437e-04,  ..., 1.3990e-05,\n",
      "           2.6124e-04, 8.4249e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.3139e-07, 2.2897e-05, 2.0723e-04,  ..., 1.2901e-05,\n",
      "           2.6724e-04, 8.4031e-05]],\n",
      "\n",
      "         [[5.3373e-07, 2.3377e-05, 2.1399e-04,  ..., 1.2647e-05,\n",
      "           2.6504e-04, 8.3203e-05]],\n",
      "\n",
      "         [[2.2104e-06, 1.7312e-05, 1.0592e-04,  ..., 9.4811e-06,\n",
      "           2.8066e-04, 5.0181e-05]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[1.7401e-06, 3.1402e-05, 4.0799e-04,  ..., 9.8536e-05,\n",
      "           3.6117e-04, 1.3315e-04]],\n",
      "\n",
      "         [[6.9754e-07, 2.6853e-05, 2.7957e-04,  ..., 1.6526e-05,\n",
      "           2.9464e-04, 9.5752e-05]],\n",
      "\n",
      "         [[7.5799e-07, 2.6606e-05, 2.8222e-04,  ..., 1.6535e-05,\n",
      "           2.9708e-04, 9.6574e-05]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[5.9358e-07, 2.3072e-05, 2.5998e-04,  ..., 1.4910e-05,\n",
      "           3.0269e-04, 9.5642e-05]],\n",
      "\n",
      "         [[6.0464e-07, 2.3727e-05, 2.6768e-04,  ..., 1.4620e-05,\n",
      "           3.0132e-04, 9.4991e-05]],\n",
      "\n",
      "         [[1.9648e-06, 1.9369e-05, 1.3575e-04,  ..., 1.0587e-05,\n",
      "           3.2189e-04, 5.2765e-05]]],\n",
      "\n",
      "\n",
      "        [[[1.8183e-06, 3.2571e-05, 4.2880e-04,  ..., 1.0308e-04,\n",
      "           3.7470e-04, 1.3831e-04]],\n",
      "\n",
      "         [[7.4557e-07, 2.8437e-05, 2.9360e-04,  ..., 1.7338e-05,\n",
      "           3.0603e-04, 9.9775e-05]],\n",
      "\n",
      "         [[8.0261e-07, 2.7964e-05, 2.9702e-04,  ..., 1.7452e-05,\n",
      "           3.0894e-04, 1.0071e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[7.2156e-07, 2.3847e-05, 2.7319e-04,  ..., 1.5724e-05,\n",
      "           3.1543e-04, 9.9863e-05]],\n",
      "\n",
      "         [[7.1458e-07, 2.4423e-05, 2.8125e-04,  ..., 1.5426e-05,\n",
      "           3.1356e-04, 9.9200e-05]],\n",
      "\n",
      "         [[2.0746e-06, 1.9952e-05, 1.4292e-04,  ..., 1.1105e-05,\n",
      "           3.3457e-04, 5.4465e-05]]],\n",
      "\n",
      "\n",
      "        [[[1.9631e-06, 3.4965e-05, 4.6223e-04,  ..., 1.1076e-04,\n",
      "           3.9801e-04, 1.4723e-04]],\n",
      "\n",
      "         [[8.0517e-07, 3.0649e-05, 3.1885e-04,  ..., 1.8846e-05,\n",
      "           3.2722e-04, 1.0685e-04]],\n",
      "\n",
      "         [[8.8193e-07, 3.0067e-05, 3.2064e-04,  ..., 1.8952e-05,\n",
      "           3.2940e-04, 1.0761e-04]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[8.4224e-07, 2.5325e-05, 2.9437e-04,  ..., 1.6946e-05,\n",
      "           3.3620e-04, 1.0659e-04]],\n",
      "\n",
      "         [[8.5018e-07, 2.5870e-05, 3.0585e-04,  ..., 1.6776e-05,\n",
      "           3.3590e-04, 1.0641e-04]],\n",
      "\n",
      "         [[2.1891e-06, 2.1495e-05, 1.5481e-04,  ..., 1.1775e-05,\n",
      "           3.5604e-04, 5.8224e-05]]]], device='cuda:0')}, 37: {'step': tensor(19527.), 'exp_avg': tensor([-4.2958e-04, -3.4913e-04,  8.7061e-04,  3.6699e-04,  1.2122e-03,\n",
      "         9.5548e-04,  1.2306e-03,  1.9314e-04,  1.3656e-03,  1.0857e-03,\n",
      "        -2.0790e-05, -4.2883e-04], device='cuda:0'), 'exp_avg_sq': tensor([0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001,\n",
      "        0.0001, 0.0001, 0.0002], device='cuda:0')}}\n",
      "param_groups \t [{'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 1e-06, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37]}]\n"
     ]
    }
   ],
   "source": [
    "for var_name in optimizerTeacher.state_dict():\n",
    "    print(var_name, \"\\t\", optimizerTeacher.state_dict()[var_name])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5d140cf4-8e71-40bc-85a6-0bd4c21e602f",
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[20], line 21\u001b[0m\n\u001b[0;32m     19\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     20\u001b[0m     loss\u001b[38;5;241m=\u001b[39mP_sloss\u001b[38;5;241m*\u001b[39malpha \u001b[38;5;66;03m#los student * alpha + 0 , loss imitasi 0\u001b[39;00m\n\u001b[1;32m---> 21\u001b[0m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     23\u001b[0m optimizerStudent\u001b[38;5;241m.\u001b[39mstep()\n\u001b[0;32m     24\u001b[0m distilation_loss \u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n",
      "File \u001b[1;32m~\\TA\\lib\\site-packages\\torch\\_tensor.py:396\u001b[0m, in \u001b[0;36mTensor.backward\u001b[1;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[0;32m    387\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m    388\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[0;32m    389\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[0;32m    390\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    394\u001b[0m         create_graph\u001b[38;5;241m=\u001b[39mcreate_graph,\n\u001b[0;32m    395\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs)\n\u001b[1;32m--> 396\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\TA\\lib\\site-packages\\torch\\autograd\\__init__.py:173\u001b[0m, in \u001b[0;36mbackward\u001b[1;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[0;32m    168\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[0;32m    170\u001b[0m \u001b[38;5;66;03m# The reason we repeat same the comment below is that\u001b[39;00m\n\u001b[0;32m    171\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[0;32m    172\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[1;32m--> 173\u001b[0m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[0;32m    174\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    175\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "#Student kd\n",
    "Studloss1=[]\n",
    "for epoch in range(40):\n",
    "    params_filename = os.path.join('./Student_MSE/', 'checkpoint_%s.pth' % epoch)\n",
    "    tmp=[]\n",
    "    tmp0=[]\n",
    "    Student.train()  # ensure dropout layers are in train mode\n",
    "    Teacher.eval()\n",
    "    for batch_index, batch_data in enumerate(train_loader):\n",
    "        encoder_inputs, labels = batch_data   # encoder_inputs torch.Size([32, 307, 1, 12])  label torch.Size([32, 307, 12])\n",
    "        optimizerStudent.zero_grad()\n",
    "        P_s = Student(encoder_inputs, edge_index_data) # torch.Size([32, 307, 12]) Ps\n",
    "        with torch.no_grad():\n",
    "            P_t=Teacher(encoder_inputs, edge_index_data)# Pt\n",
    "        P_tloss=criterionStudent(P_t,labels)\n",
    "        P_sloss=criterionStudent(P_s,labels)\n",
    "        if P_sloss > P_tloss:\n",
    "            loss=P_sloss * (alpha) + criterionStudent(P_s, P_t) * (1-alpha)\n",
    "        else:\n",
    "            loss=P_sloss*alpha #los student * alpha + 0 , loss imitasi 0\n",
    "        loss.backward()\n",
    "        \n",
    "        optimizerStudent.step()\n",
    "        distilation_loss = loss.item()\n",
    "        tmp.append(distilation_loss)\n",
    "        tmp0.append(P_sloss.item())\n",
    "        global_step += 1\n",
    "        sw.add_scalar('training_loss', P_sloss, global_step)\n",
    "        if global_step % 100 == 0:\n",
    "            print('global step: %s, Training loss: %.2f, Distilation loss: %.2f, time: %.2fs' % (global_step,P_sloss,distilation_loss, time() - start_time))\n",
    "    studloss=sum(tmp0)/len(tmp0)\n",
    "    rtloss_t=sum(tmp) / len(tmp)\n",
    "    Studloss1.append(studloss)\n",
    "    TrainLoss.append(rtloss_t)\n",
    "    print(\"rata rata Trainingloss: \",studloss,\"rata rata Distilation loss\",rtloss_t,\"pada epoch:\",epoch)\n",
    "    val_loss = compute_val_loss_mstgcn(Student, val_loader, criterionStudent, masked_flag, missing_value, sw, epoch,edge_index_data)\n",
    "\n",
    "    if val_loss < best_val_loss:\n",
    "        best_epoch=epoch\n",
    "        best_val_loss = val_loss\n",
    "        checkpoint = {\n",
    "            'epoch': best_epoch,\n",
    "            'model_state': Student.state_dict(),\n",
    "            'optimizer_state': optimizerStudent.state_dict()\n",
    "        }\n",
    "        torch.save(checkpoint, params_filename)\n",
    "        print('save parameters to file: %s' % params_filename)\n",
    "\n",
    "print('best epoch:', best_epoch)\n",
    "path=os.path.join(\"./Student_MSE/\",'losshistval.npy')\n",
    "path1=os.path.join(\"./Student_MSE/\",'losshistdistilation.npy')\n",
    "path2=os.path.join(\"./Student_MSE/\",'losshisttrain.npy')\n",
    "\n",
    "np.save(path,ValLoss)\n",
    "np.save(path1,TrainLoss)\n",
    "np.save(path2,Studloss1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f334d0ba-663c-449f-b5e8-37a56d4f9390",
   "metadata": {},
   "outputs": [],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a865f0c9-f9b5-4d9a-b248-9962dd0e0f42",
   "metadata": {},
   "source": [
    "* label smoothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b823251-7c80-4d2c-a650-cbff40a597de",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch =  37\n",
    "params_filename = os.path.join('./Student_MSE/alpha 0.6/checkpoint_%s.pth' % epoch)\n",
    "loaded_checkpoint=torch.load(params_filename)\n",
    "epoch=loaded_checkpoint['epoch']\n",
    "Student.load_state_dict(loaded_checkpoint['model_state'])\n",
    "optimizerStudent.load_state_dict(loaded_checkpoint['optimizer_state'])\n",
    "Student.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb44449f-1699-4d4b-94c3-b351e783b54d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import process_time as ps\n",
    "start = ps() \n",
    "Student.train(False)  # ensure dropout layers are in evaluation mode\n",
    "with torch.no_grad():\n",
    "    test_loader_length = len(test_loader)  # nb of batch\n",
    "    tmp = []  # batch loss\n",
    "    for batch_index, batch_data in enumerate(test_loader):\n",
    "        encoder_inputs, labels = batch_data\n",
    "        outputs = Student(encoder_inputs, edge_index_data)\n",
    "        loss = criterionStudent(outputs, labels)\n",
    "        tmp.append(loss.item())\n",
    "        if batch_index % 1 == 0:\n",
    "            print('test_loss batch %s / %s, loss: %.2f ' % (batch_index + 1, test_loader_length, loss.item()))\n",
    "\n",
    "    \n",
    "    test_loss = sum(tmp) / len(tmp)\n",
    "end = ps()\n",
    "print(\"exec time : \",end-start, \"second\", \"Test loss : \",test_loss)  \n",
    "    \n",
    "# print(test_loss)\n",
    "np.save('./Student_MSE/test.npy',tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cc16a6c-12f8-4b7c-b31e-f2be33a8789e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_output = outputs[0]  # prediction\n",
    "sample_labels = labels[0] # truth\n",
    "print(sample_output.shape, sample_labels.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "098a81df-da71-44f3-a3c6-52ca7526e2a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as mdates\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca35707-3f04-4493-9f6f-741ab86a7492",
   "metadata": {},
   "outputs": [],
   "source": [
    "base = datetime.datetime(2018, 1, 1)\n",
    "dates = [base + datetime.timedelta(hours=(i)) for i in range(24)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c126983-8b66-4ef3-a7b8-093df2f8e5cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.random.choice(range(10), 24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d414e42c-c626-4c10-a074-55a6307968e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "time=['09:00','10:00','11:00','12:00','13:00','14:00','15:00','16:00','17:00','18:00','19:00','20:00']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31b19f04-9464-41e1-b31c-16a768f79e80",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import figure\n",
    "\n",
    "figure(figsize=(30,4), dpi=80)\n",
    "for i in range(2):\n",
    "    new_i = i * 12\n",
    "    plt.plot(range(0+new_i,12+new_i),sample_output[i].detach().cpu().numpy(), color = 'red')\n",
    "    plt.plot(range(0+new_i,12+new_i),sample_labels[i].cpu().numpy(), color='blue')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db023f1-c108-46f2-925c-8e4f4bb61fe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(20,10),dpi=60)\n",
    "for i in range(1):\n",
    "    new_i = i * 12\n",
    "    print(sample_labels[i].cpu().numpy())\n",
    "    ax.scatter(time,sample_labels[i].cpu().numpy(), color='blue')\n",
    "    ax.scatter(time,sample_labels[i+24].cpu().numpy(), color='magenta')\n",
    "    ax.scatter(time,sample_labels[i+36].cpu().numpy(), color='black')\n",
    "    ax.scatter(time,sample_labels[i+48].cpu().numpy(), color='cyan')\n",
    "    ax.scatter(time,sample_labels[i+60].cpu().numpy(), color='pink')\n",
    "    ax.scatter(time,sample_labels[i+72].cpu().numpy(), color='purple')\n",
    "    ax.scatter(time,sample_labels[i+84].cpu().numpy(), color='gray')\n",
    "    ax.scatter(time,sample_labels[i+96].cpu().numpy(), color='olive')\n",
    "    ax.scatter(time,sample_labels[i+108].cpu().numpy(), color='yellow')\n",
    "    ax.scatter(time,sample_labels[i+120].cpu().numpy(), color='palegreen')\n",
    "    ax.scatter(time,sample_labels[i+132].cpu().numpy(), color='peru')\n",
    "    ax.scatter(time,sample_labels[i+144].cpu().numpy(), color='teal')\n",
    "    ax.scatter(time,sample_labels[i+156].cpu().numpy(), color='cadetblue')\n",
    "    ax.scatter(time,sample_labels[i+168].cpu().numpy(), color='indigo')\n",
    "    ax.legend(['73','154','263','56','96','42','58','95','72','271','68','134','107','130'], loc=1,fontsize = 25, fancybox = True, title=' ID detektor',title_fontsize=25,bbox_to_anchor=(1.2, 1.05))\n",
    "# fig.suptitle('Visualisasi Volume kendaraan pada 15 detektor pertama (01/01/2018) ', fontsize=30)\n",
    "plt.xticks(fontsize = 30)\n",
    "plt.yticks(fontsize = 30)\n",
    "plt.xlabel('Time Stamp', fontsize=30)\n",
    "plt.ylabel('Flow', fontsize=30)\n",
    "fig.savefig(\"Volume lalulintas.png\",dpi=400,bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afcdbdcd-7146-4190-8584-85ec89580441",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67147d2b-a3e9-415c-adb4-3e711ce56820",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(adj_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed3629c6-0e25-4e39-9fb0-033531035e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[:15]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3716899f-8fa9-4c04-b1b3-de259a83ee7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(Student)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b279fd-0db6-4e18-8275-7636b08b7867",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._temporal_attention._U1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a054f770-8d43-48c2-ab06-d88c0e4adf47",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._temporal_attention._U2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7012df3-c35d-4326-9380-8e9fa1497b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._temporal_attention._U3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb3e3499-533a-42f7-bc41-38cec71cf9cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._temporal_attention._Ve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47960aff-8521-4bbc-972c-aa11d7a6ba5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._temporal_attention._be"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a13b5a-f62e-492b-b143-340e67288f6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "w1=Student._blocklist[0]._spatial_attention._W1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e961565c-d7c9-4246-9ccb-2e7aa161a65a",
   "metadata": {},
   "outputs": [],
   "source": [
    "w2=Student._blocklist[0]._spatial_attention._W2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9704a2fe-f6db-4632-8822-fca0721326a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._spatial_attention._W3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f38d683-3dcf-4e21-a265-384c40d84d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._spatial_attention._Vs[0,:12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46e6ec34-ea96-4073-88d8-623401c1583a",
   "metadata": {},
   "outputs": [],
   "source": [
    "Student._blocklist[0]._spatial_attention._bs[0,0,:12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ea60ee6-9643-4236-87d9-e770d949118e",
   "metadata": {},
   "outputs": [],
   "source": [
    " for batch_index, batch_data in enumerate(test_loader):\n",
    "        encoder_inputs, labels = batch_data\n",
    "        np.matmul(encoder_inputs.detach.cpu.numpy(),w1)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17bbb0c2-ea81-41e5-bb0c-eb0c66751d3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_inputs.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c88e57-353f-4f31-885a-445eeacba76f",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels[0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c3f4202-b6ab-409e-ac15-7260301c5644",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aae3f98c-42ea-400a-bb98-1144ef1821a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f1df6ea-74c4-4d2e-ad45-64f702fb031a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.matmul(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d35f96b5-0454-4fe4-978f-d3ff99555f00",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}