single_character_recognition_97.2%.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dhruvak001/dhruvakkaushal.github.io/blob/main/single_character_recognition_97.2%25.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0zEWf9vg98HN",
        "outputId": "352f5713-3d71-4e3e-9db7-b99c76fe883f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 583/583 [00:08<00:00, 69.54it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "1\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 71.96it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "2\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:08<00:00, 65.97it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 583/583 [00:07<00:00, 78.55it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "4\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 80.82it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "5\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 583/583 [00:06<00:00, 84.60it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "6\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 81.42it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "7\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 603/603 [00:09<00:00, 62.58it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 80.70it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "9\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 81.46it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "A\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 87.21it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "B\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 72.85it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "C\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 78.79it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "D\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 76.71it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "E\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 73.20it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "F\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 78.42it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "G\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:08<00:00, 70.67it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "H\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 72.01it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "I\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 82.70it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "J\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 593/593 [00:07<00:00, 82.25it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "K\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 590/590 [00:07<00:00, 82.24it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "L\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 74.82it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "M\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 81.78it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "N\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:08<00:00, 67.23it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "O\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 80.66it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "P\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 583/583 [00:07<00:00, 75.27it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "Q\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 90.32it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "R\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 84.13it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "S\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 583/583 [00:07<00:00, 82.31it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "T\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 76.86it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "U\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:07<00:00, 78.76it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "V\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 85.75it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "W\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:08<00:00, 71.23it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "0\n",
            "X\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 85.84it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "0\n",
            "Y\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:09<00:00, 62.96it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "0\n",
            "Z\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 573/573 [00:06<00:00, 84.94it/s] "
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "583\n",
            "573\n",
            "603\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "593\n",
            "590\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "583\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n",
            "573\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "\"\"\"import sys\n",
        "\n",
        "sys.path.append('/user/local/lib/python3.9/site-packages')\"\"\"\n",
        "\n",
        "import os\n",
        "import cv2\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "\n",
        "REBUILD_DATA =True  #keep it true for once only\n",
        "\n",
        "class Alphabets():\n",
        "  IMG_SIZE=50\n",
        "\n",
        "  training_path =\"/content/drive/MyDrive/data/training_data/\"\n",
        "\n",
        "\n",
        "  LABELS = {\"0\":0,\"1\":1,\"2\":2,\"3\":3,\"4\":4,\"5\":5,\"6\":6,\"7\":7,\"8\":8,\"9\":9,\"A\":10,\"B\":11,\"C\":12, 'D': 13, 'E': 14, 'F': 15, 'G': 16, 'H': 17, 'I': 18, 'J': 19, 'K': 20, 'L': 21, 'M': 22, 'N': 23, 'O': 24, 'P': 25, 'Q': 26, 'R': 27, 'S': 28, 'T': 29, 'U': 30, 'V': 31, 'W': 32, 'X': 33, 'Y': 34, 'Z': 35}\n",
        "  training_data =[]\n",
        "  counts = {\"0\":0,\"1\":0,\"2\":0,\"3\":0,\"4\":0,\"5\":0,\"6\":0,\"7\":0,\"8\":0,\"9\":0,\"A\":0,\"B\":0,\"C\":0, 'D': 0, 'E': 0, 'F': 0, 'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0, 'L': 0, 'M': 0, 'N': 0, 'O': 0, 'P': 0, 'Q': 0, 'R': 0, 'S': 0, 'T': 0, 'U': 0, 'V': 0, 'W': 0, 'X': 0, 'Y': 0, 'Z': 0}\n",
        "\n",
        "\n",
        "\n",
        "  def make_training_data(self):\n",
        "    for label in self.LABELS:\n",
        "      print(label)\n",
        "      training_path =self.training_path +label\n",
        "      for f in tqdm(os.listdir(training_path)):\n",
        "        try:\n",
        "          path = os.path.join(training_path, f)\n",
        "          img = cv2.imread(path , cv2.IMREAD_GRAYSCALE)\n",
        "          img = cv2.resize(img,(self.IMG_SIZE, self.IMG_SIZE))\n",
        "          self.training_data.append([np.array(img), np.eye(36)[self.LABELS[label]]])\n",
        "\n",
        "          for count in  self.counts:\n",
        "             if label == count:\n",
        "                self.counts[count]+=1\n",
        "        except Exception as e:\n",
        "          pass\n",
        "\n",
        "      np.random.shuffle(self.training_data)\n",
        "      for count in  self.counts:\n",
        "            print(self.counts.get(count))\n",
        "\n",
        "if REBUILD_DATA:\n",
        "  alphabets = Alphabets()\n",
        "  alphabets.make_training_data()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rnjgHTdDZcFY",
        "outputId": "6a71a9f6-61b4-4143-9df3-fc19574a7750"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "Net(\n",
            "  (conv1): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1))\n",
            "  (conv2): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))\n",
            "  (conv3): Conv2d(64, 128, kernel_size=(5, 5), stride=(1, 1))\n",
            "  (fc1): Linear(in_features=512, out_features=512, bias=True)\n",
            "  (fc2): Linear(in_features=512, out_features=36, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class Net(nn.Module):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.conv1 = nn.Conv2d(1,32,5)\n",
        "    self.conv2 = nn.Conv2d(32,64,5)\n",
        "    self.conv3 = nn.Conv2d(64,128,5)\n",
        "\n",
        "    x = torch.randn(50,50).view(-1,1,50,50)\n",
        "    self._to_linear = None\n",
        "    self.convs(x)\n",
        "\n",
        "    self.fc1 = nn.Linear(self._to_linear,512)\n",
        "    self.fc2 = nn.Linear(512,36)\n",
        "\n",
        "  def convs(self,x):\n",
        "    x = F.max_pool2d(F.relu(self.conv1(x)), (2,2))\n",
        "    x = F.max_pool2d(F.relu(self.conv2(x)), (2,2))\n",
        "    x = F.max_pool2d(F.relu(self.conv3(x)), (2,2))\n",
        "\n",
        "    print(x[0].shape)\n",
        "\n",
        "    if self._to_linear is None:\n",
        "      self._to_linear = x[0].shape[0]*x[0].shape[1]*x[0].shape[2]\n",
        "    return x\n",
        "\n",
        "  def forward(self,x):\n",
        "    x = self.convs(x)\n",
        "    x = x.view(-1, self._to_linear)\n",
        "    x = F.relu(self.fc1(x))\n",
        "    x = self.fc2(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "net = Net()\n",
        "\n",
        "print (net)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 90
        },
        "id": "30T-W4OViSg_",
        "outputId": "635b9ba3-80c7-4743-86ca-07a95de85b75"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-3-8422110c60bf>:6: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ../torch/csrc/utils/tensor_new.cpp:245.)\n",
            "  x = torch.Tensor([i[0] for i in alphabets.training_data]).view(-1,50,50)\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'VAL_PCT = 0.1\\nval_size = int(len(x)*VAL_PCT)\\n\\n\\nprint(val_size)\\n\\n'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "import torch.optim as optim\n",
        "\n",
        "optimizer = optim.Adam(net.parameters(), lr=0.001)\n",
        "loss_function = nn.MSELoss()\n",
        "\n",
        "x = torch.Tensor([i[0] for i in alphabets.training_data]).view(-1,50,50)\n",
        "x = x/255.0\n",
        "y = torch.Tensor([i[1] for i in alphabets.training_data])\n",
        "\n",
        "\"\"\"VAL_PCT = 0.1\n",
        "val_size = int(len(x)*VAL_PCT)\n",
        "\n",
        "\n",
        "print(val_size)\n",
        "\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j_2Nis3jpJYN",
        "outputId": "ab7cdf78-b7dd-4da3-f226-9e63021dc829"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "20745\n"
          ]
        }
      ],
      "source": [
        "train_x = x\n",
        "train_y = y\n",
        "\n",
        "print(len(train_x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XxKcqGm4phes",
        "outputId": "72bb7d9a-51fc-4ccf-cc36-928381e8c2b3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  0%|          | 0/577 [00:00<?, ?it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 3/577 [00:00<01:27,  6.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 5/577 [00:00<01:14,  7.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 7/577 [00:00<01:08,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 9/577 [00:01<01:03,  8.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 11/577 [00:01<01:03,  8.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 13/577 [00:01<01:01,  9.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 14/577 [00:01<01:04,  8.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 16/577 [00:02<01:15,  7.44it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 18/577 [00:02<01:22,  6.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 20/577 [00:02<01:23,  6.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 22/577 [00:03<01:27,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 24/577 [00:03<01:30,  6.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 26/577 [00:03<01:29,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 28/577 [00:03<01:27,  6.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▌         | 30/577 [00:04<01:26,  6.33it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 32/577 [00:04<01:30,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 34/577 [00:04<01:30,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 36/577 [00:05<01:30,  6.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 38/577 [00:05<01:26,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 40/577 [00:05<01:25,  6.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 42/577 [00:06<01:25,  6.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 44/577 [00:06<01:22,  6.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 46/577 [00:06<01:26,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 49/577 [00:07<01:10,  7.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▊         | 50/577 [00:07<01:06,  7.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▉         | 54/577 [00:07<00:59,  8.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|▉         | 55/577 [00:07<00:58,  9.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|█         | 59/577 [00:08<00:57,  9.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 61/577 [00:08<00:58,  8.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█         | 62/577 [00:08<00:56,  9.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█▏        | 65/577 [00:09<00:56,  9.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 67/577 [00:09<00:56,  9.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 69/577 [00:09<00:58,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 71/577 [00:09<00:58,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 72/577 [00:09<00:57,  8.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 76/577 [00:10<00:53,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▎        | 78/577 [00:10<00:54,  9.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▎        | 79/577 [00:10<00:53,  9.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 83/577 [00:10<00:50,  9.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▍        | 85/577 [00:11<00:53,  9.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▌        | 87/577 [00:11<00:50,  9.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 90/577 [00:11<00:52,  9.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 92/577 [00:11<00:52,  9.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▋        | 94/577 [00:12<00:51,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 97/577 [00:12<00:53,  8.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 99/577 [00:12<00:54,  8.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 101/577 [00:12<00:53,  8.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 103/577 [00:13<00:52,  9.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 105/577 [00:13<00:51,  9.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▊        | 107/577 [00:13<00:49,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 109/577 [00:13<00:52,  8.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 111/577 [00:14<00:50,  9.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|█▉        | 113/577 [00:14<00:49,  9.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 20%|█▉        | 114/577 [00:14<00:49,  9.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 118/577 [00:14<00:50,  9.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 21%|██        | 119/577 [00:14<00:49,  9.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██▏       | 123/577 [00:15<00:49,  9.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 125/577 [00:15<00:50,  9.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 127/577 [00:15<00:49,  9.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 129/577 [00:15<00:49,  9.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 131/577 [00:16<00:50,  8.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 133/577 [00:16<00:51,  8.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 135/577 [00:16<00:49,  8.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▎       | 137/577 [00:16<00:50,  8.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 139/577 [00:17<00:52,  8.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▍       | 140/577 [00:17<00:55,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 142/577 [00:17<01:03,  6.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 144/577 [00:17<01:08,  6.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▌       | 146/577 [00:18<01:11,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 148/577 [00:18<01:09,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 150/577 [00:18<01:09,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▋       | 152/577 [00:19<01:10,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 154/577 [00:19<01:07,  6.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 156/577 [00:19<01:10,  6.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 158/577 [00:20<01:08,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 160/577 [00:20<01:10,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 162/577 [00:20<01:09,  6.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 164/577 [00:21<01:10,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 166/577 [00:21<01:08,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 168/577 [00:21<01:07,  6.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 171/577 [00:22<00:56,  7.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 173/577 [00:22<00:52,  7.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|███       | 175/577 [00:22<00:47,  8.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 177/577 [00:23<00:47,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 179/577 [00:23<00:46,  8.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███▏      | 181/577 [00:23<00:46,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 183/577 [00:23<00:44,  8.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 184/577 [00:23<00:44,  8.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 187/577 [00:24<00:42,  9.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 189/577 [00:24<00:41,  9.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 191/577 [00:24<00:41,  9.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 193/577 [00:24<00:43,  8.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 195/577 [00:25<00:43,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 197/577 [00:25<00:41,  9.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 199/577 [00:25<00:41,  9.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▍      | 201/577 [00:25<00:40,  9.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▌      | 203/577 [00:25<00:40,  9.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 205/577 [00:26<00:39,  9.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 207/577 [00:26<00:40,  9.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 208/577 [00:26<00:41,  8.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 212/577 [00:26<00:39,  9.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 37%|███▋      | 213/577 [00:26<00:39,  9.33it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 217/577 [00:27<00:37,  9.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 219/577 [00:27<00:38,  9.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 221/577 [00:27<00:39,  8.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▊      | 223/577 [00:28<00:37,  9.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 39%|███▉      | 224/577 [00:28<00:37,  9.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 227/577 [00:28<00:37,  9.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 40%|███▉      | 228/577 [00:28<00:37,  9.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|████      | 232/577 [00:29<00:36,  9.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 234/577 [00:29<00:36,  9.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 236/577 [00:29<00:35,  9.48it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 238/577 [00:29<00:35,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 240/577 [00:29<00:37,  9.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 242/577 [00:30<00:37,  9.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 243/577 [00:30<00:36,  9.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 247/577 [00:30<00:36,  9.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 249/577 [00:30<00:37,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▎     | 251/577 [00:31<00:36,  8.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 253/577 [00:31<00:36,  8.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 255/577 [00:31<00:36,  8.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▍     | 257/577 [00:31<00:35,  8.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▍     | 259/577 [00:31<00:33,  9.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 261/577 [00:32<00:34,  9.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▌     | 262/577 [00:32<00:40,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 264/577 [00:32<00:43,  7.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 266/577 [00:33<00:48,  6.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▋     | 268/577 [00:33<00:48,  6.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 270/577 [00:33<00:49,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 272/577 [00:34<00:52,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 274/577 [00:34<00:50,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 276/577 [00:34<00:50,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 278/577 [00:34<00:48,  6.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▊     | 280/577 [00:35<00:47,  6.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 282/577 [00:35<00:47,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 284/577 [00:35<00:46,  6.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|████▉     | 286/577 [00:36<00:45,  6.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|████▉     | 288/577 [00:36<00:45,  6.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|█████     | 290/577 [00:36<00:45,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 292/577 [00:37<00:46,  6.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 295/577 [00:37<00:37,  7.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████▏    | 297/577 [00:37<00:34,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 299/577 [00:38<00:32,  8.44it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 301/577 [00:38<00:32,  8.49it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 52%|█████▏    | 302/577 [00:38<00:31,  8.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 306/577 [00:38<00:27,  9.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▎    | 309/577 [00:39<00:27,  9.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 311/577 [00:39<00:27,  9.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 313/577 [00:39<00:27,  9.59it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 315/577 [00:39<00:27,  9.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 317/577 [00:39<00:28,  9.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▌    | 318/577 [00:40<00:27,  9.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 322/577 [00:40<00:27,  9.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 324/577 [00:40<00:28,  9.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 56%|█████▋    | 326/577 [00:40<00:26,  9.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 328/577 [00:41<00:27,  9.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 330/577 [00:41<00:26,  9.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 332/577 [00:41<00:26,  9.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 334/577 [00:41<00:26,  9.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 336/577 [00:41<00:26,  9.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▊    | 338/577 [00:42<00:26,  8.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 340/577 [00:42<00:27,  8.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 341/577 [00:42<00:26,  8.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|█████▉    | 345/577 [00:42<00:24,  9.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|██████    | 347/577 [00:43<00:24,  9.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|██████    | 349/577 [00:43<00:24,  9.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 351/577 [00:43<00:24,  9.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 61%|██████    | 352/577 [00:43<00:24,  9.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 356/577 [00:44<00:22,  9.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 358/577 [00:44<00:24,  9.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 360/577 [00:44<00:24,  8.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 362/577 [00:44<00:24,  8.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 364/577 [00:45<00:24,  8.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 366/577 [00:45<00:24,  8.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 368/577 [00:45<00:23,  8.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 370/577 [00:45<00:22,  9.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 372/577 [00:45<00:23,  8.64it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 65%|██████▍   | 373/577 [00:46<00:23,  8.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 65%|██████▌   | 376/577 [00:46<00:21,  9.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 379/577 [00:46<00:21,  9.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 381/577 [00:46<00:20,  9.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▋   | 383/577 [00:47<00:21,  9.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 384/577 [00:47<00:21,  8.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 386/577 [00:47<00:26,  7.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 388/577 [00:47<00:27,  6.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 390/577 [00:48<00:29,  6.44it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 392/577 [00:48<00:28,  6.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 394/577 [00:48<00:28,  6.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▊   | 396/577 [00:49<00:29,  6.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 398/577 [00:49<00:28,  6.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 400/577 [00:49<00:28,  6.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|██████▉   | 402/577 [00:50<00:27,  6.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|███████   | 404/577 [00:50<00:26,  6.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|███████   | 406/577 [00:50<00:27,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 408/577 [00:51<00:26,  6.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 410/577 [00:51<00:27,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████▏  | 412/577 [00:51<00:26,  6.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 414/577 [00:52<00:26,  6.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 416/577 [00:52<00:25,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 418/577 [00:52<00:24,  6.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 421/577 [00:53<00:19,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 73%|███████▎  | 422/577 [00:53<00:19,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 426/577 [00:53<00:16,  8.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 428/577 [00:53<00:16,  8.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 430/577 [00:54<00:16,  8.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 75%|███████▍  | 431/577 [00:54<00:16,  8.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▌  | 435/577 [00:54<00:15,  9.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 437/577 [00:54<00:15,  8.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 439/577 [00:55<00:15,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▋  | 441/577 [00:55<00:15,  9.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 443/577 [00:55<00:14,  9.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 445/577 [00:55<00:14,  9.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 447/577 [00:55<00:14,  8.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 449/577 [00:56<00:14,  8.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 451/577 [00:56<00:13,  9.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▊  | 453/577 [00:56<00:13,  9.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▉  | 455/577 [00:56<00:13,  8.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 79%|███████▉  | 456/577 [00:56<00:13,  8.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 459/577 [00:57<00:12,  9.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 461/577 [00:57<00:12,  9.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 80%|████████  | 462/577 [00:57<00:12,  9.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 466/577 [00:57<00:11,  9.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 468/577 [00:58<00:12,  8.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████▏ | 470/577 [00:58<00:12,  8.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 82%|████████▏ | 471/577 [00:58<00:12,  8.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 475/577 [00:58<00:10,  9.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 82%|████████▏ | 476/577 [00:59<00:10,  9.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 480/577 [00:59<00:10,  9.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▎ | 482/577 [00:59<00:10,  9.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 84%|████████▎ | 483/577 [00:59<00:10,  9.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 487/577 [01:00<00:10,  8.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 489/577 [01:00<00:09,  9.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▌ | 491/577 [01:00<00:09,  9.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▌ | 493/577 [01:00<00:09,  9.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 495/577 [01:01<00:09,  8.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 497/577 [01:01<00:08,  9.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▋ | 499/577 [01:01<00:08,  9.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 501/577 [01:01<00:08,  8.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 503/577 [01:02<00:08,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 505/577 [01:02<00:08,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 507/577 [01:02<00:07,  9.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 509/577 [01:02<00:07,  8.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 88%|████████▊ | 510/577 [01:02<00:07,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▊ | 512/577 [01:03<00:08,  7.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 514/577 [01:03<00:09,  6.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 516/577 [01:03<00:09,  6.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|████████▉ | 518/577 [01:04<00:09,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 520/577 [01:04<00:09,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 522/577 [01:04<00:08,  6.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 524/577 [01:05<00:08,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 526/577 [01:05<00:08,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 528/577 [01:05<00:07,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 530/577 [01:06<00:07,  6.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 532/577 [01:06<00:07,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 534/577 [01:06<00:06,  6.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 536/577 [01:06<00:06,  6.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 538/577 [01:07<00:06,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▎| 540/577 [01:07<00:06,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 542/577 [01:07<00:05,  6.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 545/577 [01:08<00:04,  7.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 547/577 [01:08<00:03,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 549/577 [01:08<00:03,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 551/577 [01:09<00:03,  8.48it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 553/577 [01:09<00:02,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 96%|█████████▌| 554/577 [01:09<00:03,  7.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▋| 556/577 [01:09<00:03,  6.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 559/577 [01:10<00:02,  6.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 561/577 [01:10<00:02,  7.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 563/577 [01:10<00:01,  8.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 565/577 [01:10<00:01,  8.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 567/577 [01:11<00:01,  8.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▊| 569/577 [01:11<00:00,  9.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 571/577 [01:11<00:00,  9.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 573/577 [01:11<00:00,  9.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|█████████▉| 575/577 [01:12<00:00,  8.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 577/577 [01:12<00:00,  7.99it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  0%|          | 0/577 [00:00<?, ?it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  0%|          | 1/577 [00:00<01:03,  9.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 3/577 [00:00<01:00,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 4/577 [00:00<01:00,  9.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 5/577 [00:00<01:00,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 6/577 [00:00<01:02,  9.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 7/577 [00:00<01:03,  9.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|▏         | 8/577 [00:00<01:04,  8.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 9/577 [00:00<01:05,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 10/577 [00:01<01:04,  8.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 11/577 [00:01<01:04,  8.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 12/577 [00:01<01:04,  8.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 13/577 [00:01<01:02,  8.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 14/577 [00:01<01:02,  9.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  3%|▎         | 15/577 [00:01<01:01,  9.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  3%|▎         | 16/577 [00:01<01:00,  9.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 18/577 [00:01<00:58,  9.49it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 22/577 [00:02<00:58,  9.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 24/577 [00:02<00:58,  9.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 26/577 [00:02<00:59,  9.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 28/577 [00:03<01:01,  8.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▌         | 30/577 [00:03<01:02,  8.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 32/577 [00:03<01:05,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 34/577 [00:03<01:05,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 36/577 [00:04<01:05,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 38/577 [00:04<01:06,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 40/577 [00:04<01:03,  8.44it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 42/577 [00:04<01:01,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 44/577 [00:04<01:01,  8.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  8%|▊         | 45/577 [00:05<01:01,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 48/577 [00:05<00:57,  9.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▊         | 50/577 [00:05<01:00,  8.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▉         | 52/577 [00:05<00:57,  9.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▉         | 53/577 [00:05<00:57,  9.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|▉         | 55/577 [00:06<01:08,  7.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|▉         | 57/577 [00:06<01:16,  6.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|█         | 59/577 [00:06<01:24,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 61/577 [00:07<01:25,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 63/577 [00:07<01:24,  6.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█▏        | 65/577 [00:07<01:24,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 67/577 [00:08<01:25,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 69/577 [00:08<01:26,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 71/577 [00:08<01:25,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 73/577 [00:09<01:23,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 75/577 [00:09<01:19,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 77/577 [00:09<01:22,  6.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▎        | 79/577 [00:10<01:24,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 81/577 [00:10<01:22,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 83/577 [00:10<01:24,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▍        | 85/577 [00:11<01:23,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▌        | 88/577 [00:11<01:07,  7.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 90/577 [00:11<01:03,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 92/577 [00:12<01:01,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▋        | 94/577 [00:12<00:56,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 96/577 [00:12<00:54,  8.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 98/577 [00:12<00:57,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 100/577 [00:13<00:54,  8.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 102/577 [00:13<00:56,  8.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 18%|█▊        | 103/577 [00:13<00:56,  8.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 106/577 [00:13<00:56,  8.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▊        | 108/577 [00:14<00:53,  8.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 110/577 [00:14<00:54,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 112/577 [00:14<00:55,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|█▉        | 114/577 [00:14<00:56,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 116/577 [00:14<00:56,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 118/577 [00:15<00:55,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██        | 120/577 [00:15<00:55,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██        | 122/577 [00:15<00:54,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██▏       | 124/577 [00:15<00:54,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 126/577 [00:16<00:52,  8.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 128/577 [00:16<00:53,  8.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 130/577 [00:16<00:52,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 132/577 [00:16<00:53,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 134/577 [00:17<00:55,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▎       | 136/577 [00:17<00:52,  8.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▎       | 137/577 [00:17<00:53,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 140/577 [00:17<00:53,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 142/577 [00:18<00:51,  8.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 144/577 [00:18<00:50,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▌       | 146/577 [00:18<00:51,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 25%|██▌       | 147/577 [00:18<00:53,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 150/577 [00:19<00:51,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▋       | 152/577 [00:19<00:51,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 154/577 [00:19<00:51,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 27%|██▋       | 155/577 [00:19<00:52,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 158/577 [00:20<00:51,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 160/577 [00:20<00:51,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 162/577 [00:20<00:51,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 164/577 [00:20<00:49,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 166/577 [00:21<00:50,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 168/577 [00:21<00:49,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 169/577 [00:21<00:53,  7.64it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 171/577 [00:21<01:01,  6.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 173/577 [00:22<01:07,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|███       | 175/577 [00:22<01:10,  5.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 177/577 [00:22<01:08,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 179/577 [00:23<01:09,  5.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███▏      | 181/577 [00:23<01:09,  5.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 183/577 [00:23<01:07,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 185/577 [00:24<01:06,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 187/577 [00:24<01:06,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 189/577 [00:24<01:08,  5.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 191/577 [00:25<01:07,  5.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 193/577 [00:25<01:07,  5.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 195/577 [00:25<01:05,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 197/577 [00:26<01:06,  5.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 199/577 [00:26<01:04,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▌      | 202/577 [00:27<00:53,  6.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▌      | 204/577 [00:27<00:48,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 206/577 [00:27<00:45,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 208/577 [00:27<00:44,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▋      | 210/577 [00:28<00:43,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 212/577 [00:28<00:43,  8.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 214/577 [00:28<00:44,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 216/577 [00:28<00:43,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 218/577 [00:28<00:42,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 220/577 [00:29<00:41,  8.52it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 222/577 [00:29<00:42,  8.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 224/577 [00:29<00:41,  8.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 226/577 [00:29<00:41,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 228/577 [00:30<00:41,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 230/577 [00:30<00:41,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|████      | 232/577 [00:30<00:40,  8.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 234/577 [00:30<00:40,  8.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 236/577 [00:31<00:40,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 238/577 [00:31<00:39,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 240/577 [00:31<00:39,  8.59it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 242/577 [00:31<00:39,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 244/577 [00:32<00:37,  8.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 246/577 [00:32<00:37,  8.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 248/577 [00:32<00:38,  8.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 250/577 [00:32<00:37,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▎     | 252/577 [00:32<00:36,  9.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 254/577 [00:33<00:36,  8.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 256/577 [00:33<00:37,  8.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▍     | 257/577 [00:33<00:38,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 260/577 [00:33<00:37,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 262/577 [00:34<00:37,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 264/577 [00:34<00:37,  8.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 266/577 [00:34<00:35,  8.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▋     | 268/577 [00:34<00:37,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 270/577 [00:35<00:37,  8.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 47%|████▋     | 271/577 [00:35<00:37,  8.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 274/577 [00:35<00:36,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 276/577 [00:35<00:35,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 278/577 [00:36<00:35,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▊     | 280/577 [00:36<00:35,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 282/577 [00:36<00:34,  8.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 49%|████▉     | 283/577 [00:36<00:35,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 285/577 [00:37<00:41,  7.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|████▉     | 287/577 [00:37<00:44,  6.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|█████     | 289/577 [00:37<00:45,  6.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|█████     | 291/577 [00:38<00:46,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 293/577 [00:38<00:46,  6.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 295/577 [00:38<00:45,  6.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████▏    | 297/577 [00:38<00:46,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 299/577 [00:39<00:47,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 301/577 [00:39<00:48,  5.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 303/577 [00:40<00:47,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 305/577 [00:40<00:48,  5.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 307/577 [00:40<00:46,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▎    | 309/577 [00:41<00:45,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 311/577 [00:41<00:45,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 313/577 [00:41<00:46,  5.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 315/577 [00:42<00:43,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▌    | 318/577 [00:42<00:36,  7.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▌    | 320/577 [00:42<00:32,  7.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 56%|█████▌    | 321/577 [00:42<00:32,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 324/577 [00:43<00:30,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▋    | 326/577 [00:43<00:28,  8.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 328/577 [00:43<00:27,  9.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 330/577 [00:43<00:28,  8.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 332/577 [00:44<00:28,  8.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 334/577 [00:44<00:28,  8.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 336/577 [00:44<00:28,  8.44it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▊    | 338/577 [00:44<00:28,  8.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 340/577 [00:45<00:28,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 342/577 [00:45<00:27,  8.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|█████▉    | 344/577 [00:45<00:27,  8.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|█████▉    | 346/577 [00:45<00:26,  8.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|██████    | 348/577 [00:46<00:27,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 350/577 [00:46<00:26,  8.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 352/577 [00:46<00:25,  8.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████▏   | 354/577 [00:46<00:25,  8.60it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 356/577 [00:46<00:26,  8.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 358/577 [00:47<00:26,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 360/577 [00:47<00:25,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 362/577 [00:47<00:25,  8.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 364/577 [00:47<00:25,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 366/577 [00:48<00:25,  8.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 368/577 [00:48<00:24,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 370/577 [00:48<00:23,  8.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 372/577 [00:48<00:23,  8.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 65%|██████▍   | 373/577 [00:48<00:24,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 65%|██████▌   | 376/577 [00:49<00:23,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 378/577 [00:49<00:24,  8.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 380/577 [00:49<00:24,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 382/577 [00:50<00:23,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 384/577 [00:50<00:23,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 386/577 [00:50<00:23,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 388/577 [00:50<00:23,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 390/577 [00:51<00:23,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 68%|██████▊   | 391/577 [00:51<00:23,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 394/577 [00:51<00:22,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▊   | 396/577 [00:51<00:21,  8.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 398/577 [00:52<00:21,  8.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 400/577 [00:52<00:20,  8.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 69%|██████▉   | 401/577 [00:52<00:22,  7.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|██████▉   | 403/577 [00:52<00:27,  6.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|███████   | 405/577 [00:53<00:27,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 407/577 [00:53<00:26,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 409/577 [00:53<00:27,  6.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 411/577 [00:54<00:27,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 413/577 [00:54<00:27,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 415/577 [00:54<00:27,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 417/577 [00:55<00:26,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 419/577 [00:55<00:26,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 421/577 [00:55<00:26,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 423/577 [00:56<00:25,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▎  | 425/577 [00:56<00:25,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 427/577 [00:56<00:25,  5.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 429/577 [00:57<00:25,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 431/577 [00:57<00:24,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▌  | 433/577 [00:57<00:24,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 436/577 [00:58<00:19,  7.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 438/577 [00:58<00:17,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 76%|███████▌  | 439/577 [00:58<00:18,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 442/577 [00:58<00:16,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 444/577 [00:59<00:16,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 446/577 [00:59<00:16,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 448/577 [00:59<00:15,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 450/577 [00:59<00:15,  8.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 452/577 [01:00<00:14,  8.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▊  | 454/577 [01:00<00:14,  8.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 79%|███████▉  | 455/577 [01:00<00:14,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▉  | 458/577 [01:00<00:13,  8.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 460/577 [01:01<00:13,  8.48it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|████████  | 462/577 [01:01<00:13,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|████████  | 464/577 [01:01<00:13,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 466/577 [01:01<00:13,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 468/577 [01:02<00:13,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████▏ | 470/577 [01:02<00:12,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 472/577 [01:02<00:12,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 474/577 [01:02<00:12,  8.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 476/577 [01:03<00:11,  8.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 478/577 [01:03<00:11,  8.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 480/577 [01:03<00:11,  8.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▎ | 482/577 [01:03<00:11,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 484/577 [01:03<00:10,  8.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 486/577 [01:04<00:10,  8.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 488/577 [01:04<00:10,  8.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 490/577 [01:04<00:09,  8.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▌ | 492/577 [01:04<00:09,  8.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 494/577 [01:05<00:09,  9.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 496/577 [01:05<00:09,  8.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▋ | 498/577 [01:05<00:09,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 500/577 [01:05<00:08,  8.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 502/577 [01:06<00:08,  8.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 504/577 [01:06<00:08,  9.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 506/577 [01:06<00:08,  8.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 508/577 [01:06<00:08,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 510/577 [01:06<00:08,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▊ | 512/577 [01:07<00:07,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 514/577 [01:07<00:07,  8.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 516/577 [01:07<00:07,  8.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|████████▉ | 518/577 [01:07<00:07,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 90%|████████▉ | 519/577 [01:08<00:07,  7.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 521/577 [01:08<00:08,  6.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 523/577 [01:08<00:08,  6.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 525/577 [01:09<00:08,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████▏| 527/577 [01:09<00:08,  5.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 529/577 [01:09<00:08,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 531/577 [01:10<00:08,  5.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 533/577 [01:10<00:07,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 535/577 [01:10<00:06,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 537/577 [01:11<00:06,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 539/577 [01:11<00:06,  5.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 541/577 [01:11<00:05,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 543/577 [01:12<00:05,  5.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 545/577 [01:12<00:05,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 547/577 [01:12<00:05,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 549/577 [01:13<00:04,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 551/577 [01:13<00:04,  6.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 554/577 [01:13<00:03,  7.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▋| 556/577 [01:14<00:02,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 97%|█████████▋| 557/577 [01:14<00:02,  7.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 560/577 [01:14<00:02,  8.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 562/577 [01:14<00:01,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 564/577 [01:15<00:01,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 566/577 [01:15<00:01,  8.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 568/577 [01:15<00:01,  8.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 570/577 [01:15<00:00,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 572/577 [01:16<00:00,  8.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 99%|█████████▉| 573/577 [01:16<00:00,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|█████████▉| 576/577 [01:16<00:00,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r100%|██████████| 577/577 [01:16<00:00,  7.53it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  0%|          | 1/577 [00:00<01:12,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  0%|          | 2/577 [00:00<01:12,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 4/577 [00:00<01:11,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 5/577 [00:00<01:09,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 6/577 [00:00<01:10,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 7/577 [00:00<01:09,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|▏         | 8/577 [00:00<01:09,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 9/577 [00:01<01:10,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 10/577 [00:01<01:09,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 11/577 [00:01<01:09,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 13/577 [00:01<01:10,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 14/577 [00:01<01:10,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  3%|▎         | 15/577 [00:01<01:08,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 17/577 [00:02<01:10,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  3%|▎         | 18/577 [00:02<01:09,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  3%|▎         | 19/577 [00:02<01:09,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  3%|▎         | 20/577 [00:02<01:10,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  4%|▎         | 21/577 [00:02<01:07,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  4%|▍         | 22/577 [00:02<01:09,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  4%|▍         | 23/577 [00:02<01:08,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  4%|▍         | 24/577 [00:02<01:08,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  4%|▍         | 25/577 [00:03<01:08,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  5%|▍         | 26/577 [00:03<01:09,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  5%|▍         | 27/577 [00:03<01:08,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  5%|▍         | 28/577 [00:03<01:08,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  5%|▌         | 29/577 [00:03<01:07,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  5%|▌         | 30/577 [00:03<01:06,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 32/577 [00:03<01:08,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  6%|▌         | 33/577 [00:04<01:08,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  6%|▌         | 34/577 [00:04<01:06,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  6%|▌         | 35/577 [00:04<01:04,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  6%|▌         | 36/577 [00:04<01:03,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  6%|▋         | 37/577 [00:04<01:01,  8.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 38/577 [00:04<01:00,  8.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 39/577 [00:04<01:02,  8.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 40/577 [00:04<01:02,  8.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 41/577 [00:05<01:03,  8.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 42/577 [00:05<01:05,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 44/577 [00:05<01:05,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  8%|▊         | 45/577 [00:05<01:04,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  8%|▊         | 46/577 [00:05<01:03,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  8%|▊         | 47/577 [00:05<01:03,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 49/577 [00:06<01:06,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▊         | 50/577 [00:06<01:03,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▉         | 51/577 [00:06<01:02,  8.48it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▉         | 52/577 [00:06<01:02,  8.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▉         | 53/577 [00:06<01:03,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  9%|▉         | 54/577 [00:06<01:02,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|▉         | 55/577 [00:06<01:04,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|▉         | 56/577 [00:06<01:04,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|▉         | 57/577 [00:07<01:09,  7.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|█         | 58/577 [00:07<01:16,  6.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|█         | 59/577 [00:07<01:22,  6.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|█         | 60/577 [00:07<01:22,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█         | 61/577 [00:07<01:23,  6.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█         | 62/577 [00:07<01:25,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█         | 63/577 [00:08<01:25,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█         | 64/577 [00:08<01:25,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█▏        | 65/577 [00:08<01:25,  5.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█▏        | 66/577 [00:08<01:25,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 67/577 [00:08<01:25,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 68/577 [00:08<01:28,  5.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 69/577 [00:09<01:28,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 70/577 [00:09<01:25,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 71/577 [00:09<01:25,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 12%|█▏        | 72/577 [00:09<01:23,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 13%|█▎        | 73/577 [00:09<01:21,  6.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 13%|█▎        | 74/577 [00:09<01:21,  6.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 13%|█▎        | 75/577 [00:10<01:22,  6.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 13%|█▎        | 76/577 [00:10<01:22,  6.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 13%|█▎        | 77/577 [00:10<01:22,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▎        | 78/577 [00:10<01:21,  6.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▎        | 79/577 [00:10<01:21,  6.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▍        | 80/577 [00:10<01:22,  6.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▍        | 81/577 [00:11<01:23,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▍        | 82/577 [00:11<01:25,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 14%|█▍        | 83/577 [00:11<01:25,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▍        | 84/577 [00:11<01:24,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▍        | 85/577 [00:11<01:24,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▍        | 86/577 [00:11<01:26,  5.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▌        | 87/577 [00:12<01:24,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▌        | 88/577 [00:12<01:24,  5.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▌        | 89/577 [00:12<01:26,  5.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 16%|█▌        | 90/577 [00:12<01:25,  5.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 16%|█▌        | 91/577 [00:12<01:24,  5.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 93/577 [00:13<01:09,  6.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 16%|█▋        | 94/577 [00:13<01:06,  7.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 96/577 [00:13<01:03,  7.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 17%|█▋        | 97/577 [00:13<01:04,  7.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 99/577 [00:13<01:01,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 17%|█▋        | 100/577 [00:13<01:00,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 102/577 [00:14<01:00,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 104/577 [00:14<00:58,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 18%|█▊        | 105/577 [00:14<00:57,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 18%|█▊        | 106/577 [00:14<00:55,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 19%|█▊        | 107/577 [00:14<00:55,  8.48it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 19%|█▊        | 108/577 [00:14<00:57,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 110/577 [00:15<00:57,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 19%|█▉        | 111/577 [00:15<00:58,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|█▉        | 113/577 [00:15<00:57,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 20%|█▉        | 114/577 [00:15<00:57,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 20%|█▉        | 115/577 [00:15<01:00,  7.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 117/577 [00:16<00:57,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 20%|██        | 118/577 [00:16<00:55,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 21%|██        | 119/577 [00:16<00:56,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 21%|██        | 120/577 [00:16<00:56,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 21%|██        | 121/577 [00:16<00:56,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██▏       | 123/577 [00:16<00:56,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 21%|██▏       | 124/577 [00:16<00:55,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 22%|██▏       | 125/577 [00:17<00:54,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 22%|██▏       | 126/577 [00:17<00:53,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 22%|██▏       | 127/577 [00:17<00:54,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 129/577 [00:17<00:55,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 130/577 [00:17<00:54,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 131/577 [00:17<00:53,  8.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 132/577 [00:17<00:52,  8.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 133/577 [00:18<00:53,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 134/577 [00:18<00:53,  8.33it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 135/577 [00:18<00:54,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▎       | 136/577 [00:18<00:55,  7.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 138/577 [00:18<00:53,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▍       | 139/577 [00:18<00:53,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▍       | 140/577 [00:18<00:53,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▍       | 141/577 [00:19<00:54,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 143/577 [00:19<00:53,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 25%|██▍       | 144/577 [00:19<00:51,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 25%|██▌       | 145/577 [00:19<00:52,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 25%|██▌       | 146/577 [00:19<00:51,  8.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 25%|██▌       | 147/577 [00:19<00:53,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 26%|██▌       | 148/577 [00:19<00:54,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 150/577 [00:20<00:53,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 26%|██▌       | 151/577 [00:20<00:54,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 153/577 [00:20<00:54,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 27%|██▋       | 154/577 [00:20<00:52,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 27%|██▋       | 155/577 [00:20<00:52,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 157/577 [00:21<00:53,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 27%|██▋       | 158/577 [00:21<00:51,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 159/577 [00:21<00:51,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 160/577 [00:21<00:51,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 161/577 [00:21<00:51,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 162/577 [00:21<00:53,  7.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 164/577 [00:21<00:52,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 166/577 [00:22<00:52,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 167/577 [00:22<00:50,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 168/577 [00:22<00:49,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 169/577 [00:22<00:50,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▉       | 170/577 [00:22<00:51,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 172/577 [00:22<00:50,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 30%|██▉       | 173/577 [00:23<00:58,  6.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 30%|███       | 174/577 [00:23<01:03,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 30%|███       | 175/577 [00:23<01:04,  6.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███       | 176/577 [00:23<01:06,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███       | 177/577 [00:23<01:06,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███       | 178/577 [00:23<01:07,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███       | 179/577 [00:24<01:09,  5.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███       | 180/577 [00:24<01:08,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███▏      | 181/577 [00:24<01:08,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 182/577 [00:24<01:08,  5.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 183/577 [00:24<01:07,  5.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 184/577 [00:24<01:06,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 185/577 [00:25<01:04,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 186/577 [00:25<01:06,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 32%|███▏      | 187/577 [00:25<01:06,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 33%|███▎      | 188/577 [00:25<01:05,  5.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 33%|███▎      | 189/577 [00:25<01:05,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 33%|███▎      | 190/577 [00:26<01:07,  5.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 33%|███▎      | 191/577 [00:26<01:08,  5.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 33%|███▎      | 192/577 [00:26<01:06,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 33%|███▎      | 193/577 [00:26<01:05,  5.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 34%|███▎      | 194/577 [00:26<01:06,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 34%|███▍      | 195/577 [00:26<01:07,  5.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 34%|███▍      | 196/577 [00:27<01:08,  5.60it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 34%|███▍      | 197/577 [00:27<01:10,  5.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 34%|███▍      | 198/577 [00:27<01:08,  5.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 34%|███▍      | 199/577 [00:27<01:05,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 35%|███▍      | 200/577 [00:27<01:03,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 35%|███▍      | 201/577 [00:27<01:02,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 35%|███▌      | 202/577 [00:28<01:03,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 35%|███▌      | 203/577 [00:28<01:04,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 35%|███▌      | 204/577 [00:28<01:06,  5.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 205/577 [00:28<01:03,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 206/577 [00:28<01:02,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 207/577 [00:28<01:03,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 208/577 [00:29<01:06,  5.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 209/577 [00:29<01:02,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 211/577 [00:29<00:53,  6.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 37%|███▋      | 212/577 [00:29<00:50,  7.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 37%|███▋      | 213/577 [00:29<00:49,  7.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 215/577 [00:30<00:47,  7.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 37%|███▋      | 216/577 [00:30<00:45,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 38%|███▊      | 217/577 [00:30<00:44,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 219/577 [00:30<00:45,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 38%|███▊      | 220/577 [00:30<00:44,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 38%|███▊      | 221/577 [00:30<00:44,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 38%|███▊      | 222/577 [00:30<00:44,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 39%|███▊      | 223/577 [00:31<00:45,  7.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 39%|███▉      | 224/577 [00:31<00:46,  7.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 226/577 [00:31<00:44,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 39%|███▉      | 227/577 [00:31<00:44,  7.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 40%|███▉      | 228/577 [00:31<00:44,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 40%|███▉      | 229/577 [00:31<00:45,  7.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|████      | 231/577 [00:32<00:45,  7.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|████      | 233/577 [00:32<00:44,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 41%|████      | 234/577 [00:32<00:43,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 41%|████      | 235/577 [00:32<00:44,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 237/577 [00:32<00:42,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 41%|████      | 238/577 [00:32<00:41,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 41%|████▏     | 239/577 [00:33<00:43,  7.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 241/577 [00:33<00:42,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 242/577 [00:33<00:41,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 243/577 [00:33<00:41,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 244/577 [00:33<00:41,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 245/577 [00:33<00:42,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 43%|████▎     | 246/577 [00:34<00:42,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 43%|████▎     | 247/577 [00:34<00:42,  7.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 43%|████▎     | 248/577 [00:34<00:42,  7.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 250/577 [00:34<00:41,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 44%|████▎     | 251/577 [00:34<00:43,  7.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 253/577 [00:34<00:43,  7.52it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 44%|████▍     | 254/577 [00:35<00:42,  7.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 256/577 [00:35<00:42,  7.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▍     | 257/577 [00:35<00:40,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▍     | 258/577 [00:35<00:40,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▍     | 259/577 [00:35<00:40,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 261/577 [00:35<00:40,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▌     | 262/577 [00:36<00:39,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▌     | 263/577 [00:36<00:39,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▌     | 264/577 [00:36<00:39,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▌     | 265/577 [00:36<00:38,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▌     | 266/577 [00:36<00:37,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▋     | 267/577 [00:36<00:39,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▋     | 268/577 [00:36<00:39,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 270/577 [00:37<00:37,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 47%|████▋     | 271/577 [00:37<00:40,  7.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 47%|████▋     | 272/577 [00:37<00:40,  7.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 274/577 [00:37<00:39,  7.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 48%|████▊     | 275/577 [00:37<00:38,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 48%|████▊     | 276/577 [00:37<00:36,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 48%|████▊     | 277/577 [00:38<00:37,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 48%|████▊     | 278/577 [00:38<00:38,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 48%|████▊     | 279/577 [00:38<00:39,  7.59it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▊     | 281/577 [00:38<00:36,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 49%|████▉     | 282/577 [00:38<00:35,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 49%|████▉     | 283/577 [00:38<00:34,  8.44it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 49%|████▉     | 284/577 [00:38<00:35,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 49%|████▉     | 285/577 [00:38<00:35,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|████▉     | 286/577 [00:39<00:35,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|████▉     | 287/577 [00:39<00:41,  7.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|████▉     | 288/577 [00:39<00:44,  6.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|█████     | 289/577 [00:39<00:46,  6.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|█████     | 290/577 [00:39<00:45,  6.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|█████     | 291/577 [00:39<00:45,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████     | 292/577 [00:40<00:44,  6.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████     | 293/577 [00:40<00:45,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████     | 294/577 [00:40<00:47,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████     | 295/577 [00:40<00:47,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████▏    | 296/577 [00:40<00:47,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████▏    | 297/577 [00:40<00:47,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 52%|█████▏    | 298/577 [00:41<00:47,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 52%|█████▏    | 299/577 [00:41<00:49,  5.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 52%|█████▏    | 300/577 [00:41<00:48,  5.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 52%|█████▏    | 301/577 [00:41<00:48,  5.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 52%|█████▏    | 302/577 [00:41<00:47,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 53%|█████▎    | 303/577 [00:42<00:47,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 53%|█████▎    | 304/577 [00:42<00:48,  5.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 53%|█████▎    | 305/577 [00:42<00:48,  5.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 53%|█████▎    | 306/577 [00:42<00:49,  5.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 53%|█████▎    | 307/577 [00:42<00:48,  5.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 53%|█████▎    | 308/577 [00:42<00:47,  5.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 54%|█████▎    | 309/577 [00:43<00:46,  5.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 54%|█████▎    | 310/577 [00:43<00:47,  5.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 54%|█████▍    | 311/577 [00:43<00:47,  5.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 54%|█████▍    | 312/577 [00:43<00:45,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 54%|█████▍    | 313/577 [00:43<00:45,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 54%|█████▍    | 314/577 [00:43<00:44,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▍    | 315/577 [00:44<00:44,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▍    | 316/577 [00:44<00:45,  5.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▍    | 317/577 [00:44<00:46,  5.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▌    | 318/577 [00:44<00:46,  5.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▌    | 319/577 [00:44<00:44,  5.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 55%|█████▌    | 320/577 [00:45<00:43,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 56%|█████▌    | 321/577 [00:45<00:42,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 56%|█████▌    | 322/577 [00:45<00:43,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 56%|█████▌    | 323/577 [00:45<00:40,  6.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▋    | 325/577 [00:45<00:36,  7.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 56%|█████▋    | 326/577 [00:45<00:33,  7.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 57%|█████▋    | 327/577 [00:45<00:32,  7.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 57%|█████▋    | 328/577 [00:46<00:32,  7.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 57%|█████▋    | 329/577 [00:46<00:31,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 57%|█████▋    | 330/577 [00:46<00:31,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 57%|█████▋    | 331/577 [00:46<00:31,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 333/577 [00:46<00:31,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 58%|█████▊    | 334/577 [00:46<00:31,  7.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 336/577 [00:47<00:30,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 58%|█████▊    | 337/577 [00:47<00:29,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 339/577 [00:47<00:30,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 340/577 [00:47<00:29,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 341/577 [00:47<00:30,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 342/577 [00:47<00:29,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 343/577 [00:48<00:28,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 60%|█████▉    | 344/577 [00:48<00:28,  8.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 60%|█████▉    | 345/577 [00:48<00:28,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|██████    | 347/577 [00:48<00:28,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 60%|██████    | 348/577 [00:48<00:28,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 350/577 [00:48<00:28,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 61%|██████    | 351/577 [00:49<00:28,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 61%|██████    | 352/577 [00:49<00:28,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 61%|██████    | 353/577 [00:49<00:28,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 61%|██████▏   | 354/577 [00:49<00:28,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 62%|██████▏   | 355/577 [00:49<00:27,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 62%|██████▏   | 356/577 [00:49<00:27,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 62%|██████▏   | 357/577 [00:49<00:29,  7.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 359/577 [00:50<00:27,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 62%|██████▏   | 360/577 [00:50<00:27,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 63%|██████▎   | 361/577 [00:50<00:27,  7.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 63%|██████▎   | 362/577 [00:50<00:27,  7.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 364/577 [00:50<00:27,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 63%|██████▎   | 365/577 [00:50<00:27,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 63%|██████▎   | 366/577 [00:50<00:26,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▎   | 367/577 [00:51<00:26,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 368/577 [00:51<00:25,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 369/577 [00:51<00:25,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 370/577 [00:51<00:25,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 371/577 [00:51<00:26,  7.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 372/577 [00:51<00:26,  7.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 65%|██████▍   | 374/577 [00:51<00:26,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 65%|██████▍   | 375/577 [00:52<00:25,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 65%|██████▌   | 376/577 [00:52<00:24,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 378/577 [00:52<00:25,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 66%|██████▌   | 379/577 [00:52<00:26,  7.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 381/577 [00:52<00:25,  7.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▋   | 383/577 [00:53<00:25,  7.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 384/577 [00:53<00:24,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 385/577 [00:53<00:24,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 386/577 [00:53<00:23,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 387/577 [00:53<00:23,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 388/577 [00:53<00:24,  7.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 390/577 [00:54<00:23,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 68%|██████▊   | 391/577 [00:54<00:23,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 68%|██████▊   | 392/577 [00:54<00:24,  7.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 394/577 [00:54<00:22,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 68%|██████▊   | 395/577 [00:54<00:22,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 69%|██████▊   | 396/577 [00:54<00:22,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 69%|██████▉   | 397/577 [00:54<00:23,  7.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 399/577 [00:55<00:22,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 69%|██████▉   | 400/577 [00:55<00:22,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 69%|██████▉   | 401/577 [00:55<00:22,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|██████▉   | 402/577 [00:55<00:23,  7.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|██████▉   | 403/577 [00:55<00:25,  6.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|███████   | 404/577 [00:55<00:27,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|███████   | 405/577 [00:56<00:27,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|███████   | 406/577 [00:56<00:26,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████   | 407/577 [00:56<00:26,  6.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████   | 408/577 [00:56<00:26,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████   | 409/577 [00:56<00:27,  6.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████   | 410/577 [00:56<00:26,  6.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████   | 411/577 [00:57<00:26,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████▏  | 412/577 [00:57<00:26,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 72%|███████▏  | 413/577 [00:57<00:27,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 72%|███████▏  | 414/577 [00:57<00:27,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 72%|███████▏  | 415/577 [00:57<00:27,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 72%|███████▏  | 416/577 [00:57<00:26,  5.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 418/577 [00:58<00:27,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 420/577 [00:58<00:26,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 422/577 [00:58<00:26,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 424/577 [00:59<00:25,  6.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 426/577 [00:59<00:25,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 428/577 [00:59<00:25,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 430/577 [01:00<00:23,  6.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 432/577 [01:00<00:24,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▌  | 434/577 [01:00<00:24,  5.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 436/577 [01:01<00:24,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 438/577 [01:01<00:23,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▋  | 441/577 [01:02<00:20,  6.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 77%|███████▋  | 442/577 [01:02<00:19,  6.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 445/577 [01:02<00:17,  7.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 447/577 [01:02<00:16,  7.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 449/577 [01:03<00:16,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 78%|███████▊  | 450/577 [01:03<00:16,  7.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▊  | 453/577 [01:03<00:15,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▉  | 455/577 [01:03<00:15,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 79%|███████▉  | 456/577 [01:03<00:15,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 459/577 [01:04<00:14,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 461/577 [01:04<00:14,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|████████  | 463/577 [01:04<00:13,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 465/577 [01:05<00:13,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 467/577 [01:05<00:13,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████▏ | 469/577 [01:05<00:13,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 471/577 [01:05<00:13,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 473/577 [01:06<00:12,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 82%|████████▏ | 474/577 [01:06<00:12,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 477/577 [01:06<00:12,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 479/577 [01:06<00:12,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 83%|████████▎ | 480/577 [01:06<00:11,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▎ | 483/577 [01:07<00:11,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 84%|████████▍ | 484/577 [01:07<00:11,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 487/577 [01:07<00:11,  7.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 489/577 [01:08<00:11,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 85%|████████▍ | 490/577 [01:08<00:10,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▌ | 493/577 [01:08<00:10,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 495/577 [01:08<00:09,  8.42it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 497/577 [01:09<00:09,  8.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▋ | 499/577 [01:09<00:09,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 87%|████████▋ | 500/577 [01:09<00:09,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 503/577 [01:09<00:09,  7.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 505/577 [01:10<00:08,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 507/577 [01:10<00:08,  8.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 509/577 [01:10<00:08,  8.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 88%|████████▊ | 510/577 [01:10<00:08,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 513/577 [01:11<00:07,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 515/577 [01:11<00:07,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 89%|████████▉ | 516/577 [01:11<00:07,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|████████▉ | 519/577 [01:11<00:07,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 90%|█████████ | 520/577 [01:11<00:07,  7.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 522/577 [01:12<00:08,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 524/577 [01:12<00:08,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 526/577 [01:13<00:08,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 528/577 [01:13<00:08,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 530/577 [01:13<00:07,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 532/577 [01:14<00:07,  6.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 534/577 [01:14<00:07,  5.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 536/577 [01:14<00:07,  5.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 538/577 [01:15<00:06,  5.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▎| 540/577 [01:15<00:06,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 542/577 [01:15<00:05,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 544/577 [01:16<00:05,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 546/577 [01:16<00:05,  5.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 548/577 [01:16<00:04,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 550/577 [01:17<00:04,  5.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 552/577 [01:17<00:04,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 554/577 [01:17<00:04,  5.49it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▋| 556/577 [01:18<00:03,  5.49it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 558/577 [01:18<00:03,  5.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 560/577 [01:18<00:02,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 562/577 [01:19<00:02,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 564/577 [01:19<00:02,  6.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 566/577 [01:19<00:01,  6.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 568/577 [01:20<00:01,  6.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 570/577 [01:20<00:01,  5.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 572/577 [01:20<00:00,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 574/577 [01:21<00:00,  5.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 577/577 [01:21<00:00,  7.07it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  0%|          | 1/577 [00:00<01:41,  5.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 3/577 [00:00<01:38,  5.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 5/577 [00:00<01:36,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 7/577 [00:01<01:36,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 9/577 [00:01<01:33,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 11/577 [00:01<01:33,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 13/577 [00:02<01:34,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 16/577 [00:02<01:21,  6.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 18/577 [00:02<01:16,  7.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 20/577 [00:03<01:12,  7.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 22/577 [00:03<01:10,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 24/577 [00:03<01:05,  8.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 26/577 [00:03<01:05,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 28/577 [00:04<01:07,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▌         | 30/577 [00:04<01:07,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 32/577 [00:04<01:07,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 34/577 [00:04<01:07,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 36/577 [00:05<01:08,  7.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 38/577 [00:05<01:05,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 40/577 [00:05<01:06,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 42/577 [00:05<01:07,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 44/577 [00:06<01:05,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 46/577 [00:06<01:05,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  8%|▊         | 47/577 [00:06<01:06,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 49/577 [00:06<01:19,  6.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▉         | 51/577 [00:07<01:25,  6.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▉         | 53/577 [00:07<01:27,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|▉         | 55/577 [00:07<01:24,  6.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|▉         | 57/577 [00:08<01:24,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|█         | 59/577 [00:08<01:21,  6.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 61/577 [00:08<01:24,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 63/577 [00:09<01:25,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█▏        | 65/577 [00:09<01:23,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 67/577 [00:09<01:24,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 69/577 [00:10<01:24,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 71/577 [00:10<01:23,  6.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 73/577 [00:10<01:25,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 75/577 [00:11<01:21,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 77/577 [00:11<01:20,  6.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▎        | 79/577 [00:11<01:26,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 81/577 [00:12<01:25,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 83/577 [00:12<01:23,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▍        | 85/577 [00:12<01:22,  5.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▌        | 88/577 [00:13<01:09,  7.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 90/577 [00:13<01:04,  7.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 92/577 [00:13<01:02,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▋        | 94/577 [00:14<01:02,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 96/577 [00:14<01:02,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 98/577 [00:14<00:58,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 100/577 [00:14<00:58,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 102/577 [00:14<00:57,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 18%|█▊        | 103/577 [00:15<00:59,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 106/577 [00:15<00:57,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▊        | 108/577 [00:15<00:55,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 19%|█▉        | 109/577 [00:15<00:58,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 112/577 [00:16<00:59,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 20%|█▉        | 113/577 [00:16<01:01,  7.54it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 116/577 [00:16<00:58,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 118/577 [00:17<00:58,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 21%|██        | 119/577 [00:17<00:58,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██        | 122/577 [00:17<00:56,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██▏       | 124/577 [00:17<00:56,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 22%|██▏       | 125/577 [00:17<00:59,  7.64it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 128/577 [00:18<00:57,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 130/577 [00:18<00:55,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 132/577 [00:18<00:53,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 134/577 [00:19<00:53,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▎       | 136/577 [00:19<00:54,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 138/577 [00:19<00:54,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 140/577 [00:19<00:52,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 142/577 [00:20<00:52,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 25%|██▍       | 143/577 [00:20<00:52,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▌       | 146/577 [00:20<00:54,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 148/577 [00:20<00:52,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 150/577 [00:21<00:52,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 26%|██▌       | 151/577 [00:21<00:52,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 154/577 [00:21<00:53,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 156/577 [00:21<00:51,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 158/577 [00:22<00:51,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 160/577 [00:22<00:51,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 28%|██▊       | 161/577 [00:22<00:55,  7.49it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 164/577 [00:22<00:53,  7.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▊       | 165/577 [00:22<00:58,  7.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 167/577 [00:23<01:03,  6.43it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 169/577 [00:23<01:08,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 171/577 [00:24<01:07,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 173/577 [00:24<01:04,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|███       | 175/577 [00:24<01:06,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 177/577 [00:24<01:05,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 179/577 [00:25<01:05,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███▏      | 181/577 [00:25<01:09,  5.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 183/577 [00:26<01:05,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 185/577 [00:26<01:04,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 187/577 [00:26<01:05,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 189/577 [00:26<01:04,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 191/577 [00:27<01:04,  6.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 193/577 [00:27<01:01,  6.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 195/577 [00:27<01:02,  6.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 197/577 [00:28<01:02,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 199/577 [00:28<01:07,  5.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▍      | 201/577 [00:29<01:02,  6.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▌      | 204/577 [00:29<00:52,  7.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 36%|███▌      | 205/577 [00:29<00:50,  7.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 208/577 [00:29<00:45,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▋      | 210/577 [00:30<00:46,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 212/577 [00:30<00:44,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 214/577 [00:30<00:44,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 216/577 [00:30<00:43,  8.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 218/577 [00:31<00:43,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 38%|███▊      | 219/577 [00:31<00:44,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 222/577 [00:31<00:42,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 224/577 [00:31<00:41,  8.52it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 226/577 [00:32<00:42,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 228/577 [00:32<00:44,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 230/577 [00:32<00:43,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|████      | 232/577 [00:32<00:43,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 234/577 [00:33<00:43,  7.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 236/577 [00:33<00:41,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 238/577 [00:33<00:40,  8.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 240/577 [00:33<00:40,  8.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 242/577 [00:34<00:41,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 243/577 [00:34<00:42,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 246/577 [00:34<00:41,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 43%|████▎     | 247/577 [00:34<00:42,  7.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 250/577 [00:35<00:41,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▎     | 252/577 [00:35<00:39,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 254/577 [00:35<00:39,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 44%|████▍     | 255/577 [00:35<00:40,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▍     | 258/577 [00:36<00:39,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 45%|████▍     | 259/577 [00:36<00:40,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 262/577 [00:36<00:40,  7.70it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 264/577 [00:36<00:38,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 266/577 [00:37<00:38,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 46%|████▋     | 267/577 [00:37<00:39,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 270/577 [00:37<00:37,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 272/577 [00:37<00:37,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 274/577 [00:38<00:38,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 276/577 [00:38<00:38,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 278/577 [00:38<00:37,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▊     | 280/577 [00:38<00:34,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 49%|████▊     | 281/577 [00:39<00:36,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 283/577 [00:39<00:40,  7.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 285/577 [00:39<00:45,  6.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|████▉     | 287/577 [00:40<00:48,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|█████     | 289/577 [00:40<00:48,  5.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|█████     | 291/577 [00:40<00:49,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 293/577 [00:41<00:46,  6.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 295/577 [00:41<00:46,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████▏    | 297/577 [00:41<00:49,  5.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 299/577 [00:42<00:49,  5.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 301/577 [00:42<00:49,  5.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 303/577 [00:42<00:46,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 305/577 [00:43<00:43,  6.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 307/577 [00:43<00:44,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▎    | 309/577 [00:43<00:47,  5.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 311/577 [00:44<00:45,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 313/577 [00:44<00:44,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 315/577 [00:44<00:44,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 317/577 [00:45<00:44,  5.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▌    | 320/577 [00:45<00:38,  6.63it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 322/577 [00:45<00:34,  7.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 324/577 [00:46<00:31,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▋    | 326/577 [00:46<00:30,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 328/577 [00:46<00:30,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 330/577 [00:46<00:29,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 332/577 [00:47<00:29,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 334/577 [00:47<00:29,  8.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 336/577 [00:47<00:29,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▊    | 338/577 [00:47<00:28,  8.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 340/577 [00:47<00:27,  8.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 342/577 [00:48<00:28,  8.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 343/577 [00:48<00:30,  7.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|█████▉    | 346/577 [00:48<00:29,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|██████    | 348/577 [00:48<00:28,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 350/577 [00:49<00:28,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 61%|██████    | 351/577 [00:49<00:28,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 353/577 [00:49<00:29,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 356/577 [00:50<00:28,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 358/577 [00:50<00:27,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 360/577 [00:50<00:27,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 63%|██████▎   | 361/577 [00:50<00:27,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 364/577 [00:51<00:26,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 366/577 [00:51<00:25,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 368/577 [00:51<00:25,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 369/577 [00:51<00:25,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 372/577 [00:52<00:24,  8.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 65%|██████▍   | 374/577 [00:52<00:23,  8.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 65%|██████▍   | 375/577 [00:52<00:24,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 378/577 [00:52<00:24,  8.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 380/577 [00:52<00:24,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 382/577 [00:53<00:23,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 384/577 [00:53<00:23,  8.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 386/577 [00:53<00:22,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 388/577 [00:53<00:22,  8.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 67%|██████▋   | 389/577 [00:54<00:23,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 392/577 [00:54<00:22,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 394/577 [00:54<00:22,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 68%|██████▊   | 395/577 [00:54<00:23,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 398/577 [00:55<00:22,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 69%|██████▉   | 399/577 [00:55<00:23,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 401/577 [00:55<00:25,  6.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|██████▉   | 403/577 [00:56<00:27,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|███████   | 405/577 [00:56<00:27,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 407/577 [00:56<00:26,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 409/577 [00:57<00:27,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 411/577 [00:57<00:28,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 413/577 [00:57<00:27,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 415/577 [00:58<00:27,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 417/577 [00:58<00:27,  5.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 419/577 [00:58<00:26,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 421/577 [00:59<00:25,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 423/577 [00:59<00:26,  5.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▎  | 425/577 [00:59<00:25,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 427/577 [01:00<00:27,  5.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 74%|███████▍  | 428/577 [01:00<00:27,  5.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 430/577 [01:00<00:25,  5.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 432/577 [01:00<00:25,  5.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▌  | 434/577 [01:01<00:25,  5.60it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 436/577 [01:01<00:24,  5.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 439/577 [01:02<00:19,  7.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▋  | 441/577 [01:02<00:18,  7.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 443/577 [01:02<00:17,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 445/577 [01:02<00:15,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 447/577 [01:03<00:15,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 449/577 [01:03<00:16,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 451/577 [01:03<00:15,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▊  | 453/577 [01:03<00:15,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▉  | 455/577 [01:04<00:14,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▉  | 457/577 [01:04<00:15,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 459/577 [01:04<00:15,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 461/577 [01:04<00:13,  8.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|████████  | 463/577 [01:05<00:13,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 465/577 [01:05<00:14,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 467/577 [01:05<00:13,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████▏ | 469/577 [01:05<00:13,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 471/577 [01:06<00:12,  8.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 473/577 [01:06<00:12,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 475/577 [01:06<00:12,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 477/577 [01:06<00:12,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 479/577 [01:07<00:11,  8.34it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 481/577 [01:07<00:11,  8.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 84%|████████▎ | 482/577 [01:07<00:11,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 485/577 [01:07<00:11,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 487/577 [01:08<00:11,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 489/577 [01:08<00:11,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▌ | 491/577 [01:08<00:10,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 85%|████████▌ | 492/577 [01:08<00:10,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 495/577 [01:09<00:10,  8.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 497/577 [01:09<00:10,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 86%|████████▋ | 498/577 [01:09<00:10,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 501/577 [01:09<00:09,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 503/577 [01:10<00:09,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 505/577 [01:10<00:08,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 88%|████████▊ | 506/577 [01:10<00:09,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 509/577 [01:10<00:08,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▊ | 511/577 [01:11<00:08,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 513/577 [01:11<00:08,  7.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 89%|████████▉ | 514/577 [01:11<00:08,  7.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 516/577 [01:11<00:08,  7.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|████████▉ | 518/577 [01:12<00:09,  6.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 520/577 [01:12<00:09,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 522/577 [01:12<00:09,  5.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 524/577 [01:13<00:08,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 526/577 [01:13<00:08,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 528/577 [01:13<00:08,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 530/577 [01:14<00:08,  5.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 532/577 [01:14<00:07,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 534/577 [01:14<00:07,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 536/577 [01:15<00:06,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 538/577 [01:15<00:06,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▎| 540/577 [01:15<00:06,  6.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 542/577 [01:16<00:05,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 544/577 [01:16<00:05,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 546/577 [01:16<00:05,  6.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 548/577 [01:17<00:04,  5.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 550/577 [01:17<00:04,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 552/577 [01:17<00:04,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 555/577 [01:18<00:03,  6.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 96%|█████████▋| 556/577 [01:18<00:03,  6.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 559/577 [01:18<00:02,  7.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 97%|█████████▋| 560/577 [01:18<00:02,  7.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 563/577 [01:19<00:01,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 565/577 [01:19<00:01,  7.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 567/577 [01:19<00:01,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 98%|█████████▊| 568/577 [01:19<00:01,  7.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 571/577 [01:20<00:00,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 99%|█████████▉| 572/577 [01:20<00:00,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|█████████▉| 575/577 [01:20<00:00,  7.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 577/577 [01:21<00:00,  7.12it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  0%|          | 2/577 [00:00<01:12,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  1%|          | 3/577 [00:00<01:15,  7.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 5/577 [00:00<01:14,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|▏         | 8/577 [00:01<01:13,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  2%|▏         | 9/577 [00:01<01:12,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 12/577 [00:01<01:13,  7.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  2%|▏         | 14/577 [00:01<01:10,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 16/577 [00:02<01:08,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 18/577 [00:02<01:07,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  3%|▎         | 20/577 [00:02<01:07,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▍         | 22/577 [00:02<01:07,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  4%|▍         | 23/577 [00:02<01:08,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 26/577 [00:03<01:09,  7.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▍         | 28/577 [00:03<01:10,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  5%|▌         | 30/577 [00:03<01:06,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 32/577 [00:04<01:06,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 34/577 [00:04<01:05,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  6%|▌         | 36/577 [00:04<01:06,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 38/577 [00:04<01:08,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 39/577 [00:04<01:08,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  7%|▋         | 42/577 [00:05<01:08,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 44/577 [00:05<01:05,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 46/577 [00:05<01:07,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  8%|▊         | 48/577 [00:06<01:03,  8.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  8%|▊         | 49/577 [00:06<01:06,  7.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▉         | 52/577 [00:06<01:05,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  9%|▉         | 54/577 [00:06<01:06,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|▉         | 56/577 [00:07<01:03,  8.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 10%|▉         | 57/577 [00:07<01:11,  7.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|█         | 59/577 [00:07<01:16,  6.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 61/577 [00:07<01:20,  6.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█         | 63/577 [00:08<01:20,  6.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 11%|█▏        | 65/577 [00:08<01:25,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 67/577 [00:08<01:22,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 69/577 [00:09<01:20,  6.32it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 12%|█▏        | 71/577 [00:09<01:23,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 73/577 [00:09<01:23,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 75/577 [00:10<01:24,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 13%|█▎        | 77/577 [00:10<01:26,  5.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▎        | 79/577 [00:10<01:24,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 81/577 [00:11<01:22,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 14%|█▍        | 83/577 [00:11<01:24,  5.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▍        | 85/577 [00:11<01:20,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▌        | 87/577 [00:12<01:20,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 15%|█▌        | 89/577 [00:12<01:23,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 91/577 [00:12<01:24,  5.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▌        | 93/577 [00:13<01:20,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 16%|█▋        | 95/577 [00:13<01:24,  5.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 97/577 [00:13<01:19,  6.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 17%|█▋        | 100/577 [00:14<01:06,  7.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 102/577 [00:14<01:02,  7.64it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 18%|█▊        | 104/577 [00:14<00:59,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 18%|█▊        | 105/577 [00:14<01:00,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▊        | 108/577 [00:15<00:57,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 19%|█▉        | 109/577 [00:15<00:58,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 19%|█▉        | 112/577 [00:15<00:59,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|█▉        | 114/577 [00:16<00:58,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 116/577 [00:16<00:57,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 20%|██        | 118/577 [00:16<00:56,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██        | 120/577 [00:16<00:57,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██        | 122/577 [00:17<00:57,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 21%|██▏       | 124/577 [00:17<00:55,  8.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 126/577 [00:17<00:54,  8.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 22%|██▏       | 128/577 [00:17<00:55,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 22%|██▏       | 129/577 [00:17<00:57,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 23%|██▎       | 132/577 [00:18<00:57,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 23%|██▎       | 133/577 [00:18<00:56,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▎       | 136/577 [00:18<00:55,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 138/577 [00:19<00:55,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 24%|██▍       | 140/577 [00:19<00:53,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 24%|██▍       | 141/577 [00:19<00:54,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▍       | 144/577 [00:19<00:52,  8.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▌       | 146/577 [00:20<00:52,  8.16it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 148/577 [00:20<00:53,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▌       | 150/577 [00:20<00:52,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 26%|██▋       | 152/577 [00:20<00:53,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 154/577 [00:21<00:53,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 156/577 [00:21<00:51,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 27%|██▋       | 158/577 [00:21<00:50,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 160/577 [00:21<00:51,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 162/577 [00:22<00:51,  8.00it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 28%|██▊       | 164/577 [00:22<00:50,  8.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 29%|██▊       | 165/577 [00:22<00:51,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 168/577 [00:22<00:52,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 29%|██▉       | 170/577 [00:23<00:51,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|██▉       | 172/577 [00:23<00:51,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 30%|███       | 174/577 [00:23<00:50,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 176/577 [00:23<00:51,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 31%|███       | 177/577 [00:24<00:54,  7.40it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███       | 179/577 [00:24<01:02,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 31%|███▏      | 181/577 [00:24<01:03,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 183/577 [00:25<01:04,  6.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 185/577 [00:25<01:07,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 32%|███▏      | 187/577 [00:25<01:06,  5.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 189/577 [00:26<01:03,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 191/577 [00:26<01:03,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 193/577 [00:26<01:06,  5.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 195/577 [00:27<01:04,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 197/577 [00:27<01:04,  5.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 34%|███▍      | 199/577 [00:27<01:07,  5.61it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▍      | 201/577 [00:28<01:01,  6.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 35%|███▌      | 203/577 [00:28<01:01,  6.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 205/577 [00:28<01:03,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 207/577 [00:29<01:03,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 36%|███▌      | 209/577 [00:29<01:04,  5.71it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 211/577 [00:29<01:01,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 213/577 [00:30<01:01,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 37%|███▋      | 215/577 [00:30<01:01,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 218/577 [00:30<00:52,  6.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 38%|███▊      | 220/577 [00:31<00:47,  7.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 38%|███▊      | 221/577 [00:31<00:48,  7.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 224/577 [00:31<00:46,  7.64it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 39%|███▉      | 226/577 [00:31<00:43,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 228/577 [00:32<00:43,  8.09it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 230/577 [00:32<00:44,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 40%|████      | 231/577 [00:32<00:46,  7.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 234/577 [00:32<00:43,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 236/577 [00:33<00:42,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 41%|████      | 238/577 [00:33<00:43,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 240/577 [00:33<00:42,  7.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 42%|████▏     | 241/577 [00:33<00:44,  7.60it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 42%|████▏     | 244/577 [00:34<00:41,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 246/577 [00:34<00:41,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 43%|████▎     | 248/577 [00:34<00:42,  7.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 43%|████▎     | 249/577 [00:34<00:42,  7.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▎     | 252/577 [00:35<00:41,  7.74it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 44%|████▍     | 254/577 [00:35<00:39,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 44%|████▍     | 255/577 [00:35<00:40,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▍     | 258/577 [00:36<00:40,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 260/577 [00:36<00:39,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 45%|████▌     | 262/577 [00:36<00:40,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 264/577 [00:36<00:39,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▌     | 266/577 [00:37<00:38,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 46%|████▋     | 268/577 [00:37<00:38,  7.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 47%|████▋     | 269/577 [00:37<00:39,  7.75it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 272/577 [00:37<00:40,  7.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 274/577 [00:38<00:38,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 48%|████▊     | 275/577 [00:38<00:38,  7.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 48%|████▊     | 278/577 [00:38<00:38,  7.69it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▊     | 280/577 [00:38<00:38,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 282/577 [00:39<00:37,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 49%|████▉     | 284/577 [00:39<00:37,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|████▉     | 286/577 [00:39<00:37,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|████▉     | 288/577 [00:39<00:36,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 50%|█████     | 290/577 [00:40<00:36,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 292/577 [00:40<00:36,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 51%|█████     | 293/577 [00:40<00:35,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████     | 295/577 [00:40<00:38,  7.24it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 51%|█████▏    | 297/577 [00:41<00:43,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 299/577 [00:41<00:45,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 52%|█████▏    | 301/577 [00:41<00:43,  6.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 303/577 [00:42<00:46,  5.96it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 305/577 [00:42<00:59,  4.58it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 53%|█████▎    | 307/577 [00:42<00:51,  5.21it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▎    | 309/577 [00:43<00:50,  5.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 311/577 [00:43<00:47,  5.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 313/577 [00:44<00:45,  5.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 315/577 [00:44<00:42,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▍    | 317/577 [00:44<00:42,  6.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 55%|█████▌    | 319/577 [00:45<00:42,  6.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 321/577 [00:45<00:42,  6.05it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▌    | 323/577 [00:45<00:40,  6.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 56%|█████▋    | 325/577 [00:46<00:41,  6.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 327/577 [00:46<00:43,  5.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 329/577 [00:46<00:43,  5.68it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 57%|█████▋    | 331/577 [00:47<00:42,  5.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 334/577 [00:47<00:36,  6.73it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 58%|█████▊    | 336/577 [00:47<00:33,  7.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▊    | 338/577 [00:47<00:31,  7.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 59%|█████▉    | 339/577 [00:48<00:31,  7.59it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 59%|█████▉    | 342/577 [00:48<00:30,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|█████▉    | 344/577 [00:48<00:28,  8.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 60%|█████▉    | 345/577 [00:48<00:29,  7.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 60%|██████    | 348/577 [00:49<00:29,  7.86it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 60%|██████    | 349/577 [00:49<00:29,  7.65it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████    | 352/577 [00:49<00:28,  7.84it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 61%|██████▏   | 354/577 [00:50<00:28,  7.90it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 356/577 [00:50<00:27,  8.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 62%|██████▏   | 357/577 [00:50<00:29,  7.52it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 360/577 [00:50<00:26,  8.04it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 362/577 [00:51<00:26,  7.97it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 63%|██████▎   | 364/577 [00:51<00:26,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 63%|██████▎   | 365/577 [00:51<00:27,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 368/577 [00:51<00:26,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 64%|██████▍   | 369/577 [00:51<00:26,  7.78it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 64%|██████▍   | 371/577 [00:52<00:27,  7.50it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 65%|██████▍   | 373/577 [00:52<00:27,  7.33it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 65%|██████▌   | 376/577 [00:52<00:26,  7.72it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 65%|██████▌   | 377/577 [00:53<00:26,  7.56it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▌   | 380/577 [00:53<00:25,  7.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 66%|██████▌   | 381/577 [00:53<00:27,  7.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 66%|██████▋   | 383/577 [00:53<00:25,  7.47it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 386/577 [00:54<00:25,  7.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 67%|██████▋   | 388/577 [00:54<00:23,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 390/577 [00:54<00:23,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 392/577 [00:55<00:23,  7.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 68%|██████▊   | 394/577 [00:55<00:22,  8.06it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 68%|██████▊   | 395/577 [00:55<00:23,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 397/577 [00:55<00:23,  7.57it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 400/577 [00:56<00:22,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|██████▉   | 402/577 [00:56<00:22,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 70%|███████   | 404/577 [00:56<00:22,  7.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|███████   | 405/577 [00:56<00:22,  7.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 408/577 [00:57<00:21,  7.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 71%|███████   | 409/577 [00:57<00:21,  7.91it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 71%|███████   | 411/577 [00:57<00:24,  6.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 413/577 [00:57<00:27,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 415/577 [00:58<00:26,  6.17it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 72%|███████▏  | 417/577 [00:58<00:25,  6.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 419/577 [00:58<00:25,  6.25it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 421/577 [00:59<00:23,  6.51it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 73%|███████▎  | 423/577 [00:59<00:24,  6.36it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▎  | 425/577 [00:59<00:23,  6.53it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 427/577 [01:00<00:24,  6.19it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 74%|███████▍  | 429/577 [01:00<00:24,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▍  | 431/577 [01:00<00:23,  6.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▌  | 433/577 [01:01<00:23,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 75%|███████▌  | 435/577 [01:01<00:23,  5.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 437/577 [01:01<00:22,  6.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▌  | 439/577 [01:02<00:23,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 76%|███████▋  | 441/577 [01:02<00:22,  6.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 443/577 [01:02<00:21,  6.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 445/577 [01:03<00:22,  5.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 447/577 [01:03<00:20,  6.26it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 449/577 [01:03<00:21,  5.98it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 78%|███████▊  | 451/577 [01:04<00:21,  5.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▊  | 454/577 [01:04<00:17,  7.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 79%|███████▉  | 456/577 [01:04<00:16,  7.48it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 79%|███████▉  | 457/577 [01:04<00:16,  7.45it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|███████▉  | 460/577 [01:05<00:15,  7.77it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|████████  | 462/577 [01:05<00:14,  7.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 80%|████████  | 464/577 [01:05<00:14,  7.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████  | 466/577 [01:06<00:13,  8.23it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 81%|████████  | 467/577 [01:06<00:14,  7.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 81%|████████▏ | 470/577 [01:06<00:13,  8.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 472/577 [01:06<00:12,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 82%|████████▏ | 474/577 [01:07<00:12,  8.29it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 82%|████████▏ | 475/577 [01:07<00:13,  7.66it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 478/577 [01:07<00:12,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 83%|████████▎ | 480/577 [01:07<00:11,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▎ | 482/577 [01:08<00:11,  8.30it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 484/577 [01:08<00:11,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 486/577 [01:08<00:11,  8.22it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 488/577 [01:08<00:10,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▍ | 490/577 [01:09<00:10,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 85%|████████▌ | 492/577 [01:09<00:10,  7.99it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 85%|████████▌ | 493/577 [01:09<00:10,  7.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▌ | 496/577 [01:09<00:10,  7.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 86%|████████▋ | 498/577 [01:10<00:09,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 86%|████████▋ | 499/577 [01:10<00:10,  7.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 502/577 [01:10<00:09,  7.87it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 87%|████████▋ | 504/577 [01:10<00:09,  7.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 506/577 [01:11<00:08,  8.27it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 508/577 [01:11<00:08,  8.11it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 88%|████████▊ | 510/577 [01:11<00:08,  8.01it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▊ | 512/577 [01:11<00:07,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 514/577 [01:12<00:07,  8.14it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 89%|████████▉ | 516/577 [01:12<00:07,  8.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 90%|████████▉ | 517/577 [01:12<00:07,  7.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 520/577 [01:12<00:07,  8.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 90%|█████████ | 522/577 [01:13<00:06,  8.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 524/577 [01:13<00:06,  8.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 91%|█████████ | 526/577 [01:13<00:06,  8.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 528/577 [01:13<00:05,  8.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 530/577 [01:14<00:05,  8.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 92%|█████████▏| 531/577 [01:14<00:05,  8.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 533/577 [01:14<00:06,  6.89it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 535/577 [01:14<00:06,  6.39it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 537/577 [01:15<00:06,  6.49it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 93%|█████████▎| 539/577 [01:15<00:06,  6.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 541/577 [01:15<00:06,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 543/577 [01:16<00:05,  5.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 94%|█████████▍| 545/577 [01:16<00:05,  5.60it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▍| 547/577 [01:16<00:05,  5.94it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 549/577 [01:17<00:04,  5.85it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 95%|█████████▌| 551/577 [01:17<00:04,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 553/577 [01:17<00:03,  6.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 96%|█████████▌| 555/577 [01:18<00:03,  6.12it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 557/577 [01:18<00:03,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 559/577 [01:18<00:03,  5.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 97%|█████████▋| 561/577 [01:19<00:02,  6.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 563/577 [01:19<00:02,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 565/577 [01:19<00:02,  5.95it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 98%|█████████▊| 567/577 [01:20<00:01,  5.80it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▊| 569/577 [01:20<00:01,  5.92it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 571/577 [01:20<00:00,  6.08it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 574/577 [01:21<00:00,  6.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|█████████▉| 576/577 [01:21<00:00,  7.41it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r100%|██████████| 577/577 [01:21<00:00,  7.07it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 4. Loss: 0.0004992518806830049\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "BATCH_SIZE = 36\n",
        "EPOCHS = 5\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "  for i in tqdm(range(0, len(train_x), BATCH_SIZE)):\n",
        "    #print(i, i+BATCH_SIZE)\n",
        "    batch_x = train_x[i:i+BATCH_SIZE].view(-1,1,50,50)\n",
        "    batch_y = train_y[i:i+BATCH_SIZE]\n",
        "\n",
        "    net.zero_grad()\n",
        "    outputs = net(batch_x)\n",
        "    loss = loss_function(outputs, batch_y)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "print(f\"Epoch: {epoch}. Loss: {loss}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z7UrrJhRHE88",
        "outputId": "4f8ea69c-ff6e-4179-9228-35edc79808fd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.04it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "1\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.07it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "2\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "4\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.02it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "5\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.06it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "6\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.03it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "7\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.09it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "9\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.10it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "A\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.02it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "B\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.07it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "C\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "D\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.11it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "E\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.13it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "F\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.12it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "G\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.07it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "H\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.05it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "I\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.06it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "J\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "K\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "L\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.13it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "M\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.07it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "N\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.03it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "O\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.03it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "P\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.07it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "Q\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.09it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "R\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.11it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "S\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.13it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "T\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.09it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "U\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.08it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "V\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.05it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "0\n",
            "W\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.10it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "0\n",
            "X\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.09it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "0\n",
            "Y\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.11it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "0\n",
            "Z\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28/28 [00:13<00:00,  2.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n",
            "28\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "class testAlphabets():\n",
        "  IMG_SIZE=50\n",
        "\n",
        "  testing_path =\"/content/drive/MyDrive/data/testing_data/\"\n",
        "\n",
        "\n",
        "  LABELS = {\"0\":0,\"1\":1,\"2\":2,\"3\":3,\"4\":4,\"5\":5,\"6\":6,\"7\":7,\"8\":8,\"9\":9,\"A\":10,\"B\":11,\"C\":12, 'D': 13, 'E': 14, 'F': 15, 'G': 16, 'H': 17, 'I': 18, 'J': 19, 'K': 20, 'L': 21, 'M': 22, 'N': 23, 'O': 24, 'P': 25, 'Q': 26, 'R': 27, 'S': 28, 'T': 29, 'U': 30, 'V': 31, 'W': 32, 'X': 33, 'Y': 34, 'Z': 35}\n",
        "  testing_data =[]\n",
        "  counts = {\"0\":0,\"1\":0,\"2\":0,\"3\":0,\"4\":0,\"5\":0,\"6\":0,\"7\":0,\"8\":0,\"9\":0,\"A\":0,\"B\":0,\"C\":0, 'D': 0, 'E': 0, 'F': 0, 'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0, 'L': 0, 'M': 0, 'N': 0, 'O': 0, 'P': 0, 'Q': 0, 'R': 0, 'S': 0, 'T': 0, 'U': 0, 'V': 0, 'W': 0, 'X': 0, 'Y': 0, 'Z': 0}\n",
        "\n",
        "\n",
        "\n",
        "  def make_testing_data(self):\n",
        "    for label in self.LABELS:\n",
        "      print(label)\n",
        "      testing_path =self.testing_path +label\n",
        "      for f in tqdm(os.listdir(testing_path)):\n",
        "        try:\n",
        "          path = os.path.join(testing_path, f)\n",
        "          img = cv2.imread(path , cv2.IMREAD_GRAYSCALE)\n",
        "          img = cv2.resize(img,(self.IMG_SIZE, self.IMG_SIZE))\n",
        "          self.testing_data.append([np.array(img), np.eye(36)[self.LABELS[label]]])\n",
        "\n",
        "          for count in  self.counts:\n",
        "             if label == count:\n",
        "                self.counts[count]+=1\n",
        "        except Exception as e:\n",
        "          pass\n",
        "\n",
        "      np.random.shuffle(self.testing_data)\n",
        "      for count in  self.counts:\n",
        "            print(self.counts.get(count))\n",
        "\n",
        "if REBUILD_DATA:\n",
        "  testalphabets = testAlphabets()\n",
        "  testalphabets.make_testing_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zVKbXKOZHGKT",
        "outputId": "16060473-2629-4898-e47c-daac1bf2b77f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1008\n"
          ]
        }
      ],
      "source": [
        "x = torch.Tensor([i[0] for i in testalphabets.testing_data]).view(-1,50,50)\n",
        "x = x/255.0\n",
        "y = torch.Tensor([i[1] for i in testalphabets.testing_data])\n",
        "\n",
        "test_x = x\n",
        "test_y = y\n",
        "\n",
        "print(len(test_x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xNyXjTiizjgn",
        "outputId": "9418de6b-c102-425f-c583-093374c7d31b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  4%|▎         | 36/1008 [00:00<00:02, 358.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  7%|▋         | 74/1008 [00:00<00:02, 365.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 11%|█         | 111/1008 [00:00<00:02, 363.10it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 15%|█▍        | 148/1008 [00:00<00:02, 362.18it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 18%|█▊        | 185/1008 [00:00<00:02, 356.03it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 25%|██▌       | 257/1008 [00:00<00:02, 340.38it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 33%|███▎      | 330/1008 [00:00<00:01, 345.31it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 40%|███▉      | 402/1008 [00:01<00:01, 350.52it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 47%|████▋     | 476/1008 [00:01<00:01, 358.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 54%|█████▍    | 548/1008 [00:01<00:01, 351.76it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 62%|██████▏   | 621/1008 [00:01<00:01, 343.28it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 69%|██████▉   | 699/1008 [00:01<00:00, 364.37it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 77%|███████▋  | 775/1008 [00:02<00:00, 368.35it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 84%|████████▍ | 851/1008 [00:02<00:00, 366.81it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 92%|█████████▏| 926/1008 [00:02<00:00, 363.20it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 99%|█████████▉| 1000/1008 [00:02<00:00, 326.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r100%|██████████| 1008/1008 [00:02<00:00, 349.88it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "torch.Size([128, 2, 2])\n",
            "Accuracy : 0.972\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "correct = 0\n",
        "total = 0\n",
        "with torch.no_grad():\n",
        "  for i in tqdm(range(len(test_x))):\n",
        "    real_class = torch.argmax(test_y[i])\n",
        "    net_out = net(test_x[i].view(-1,1,50,50))[0]\n",
        "    predicted_class = torch.argmax(net_out)\n",
        "    if predicted_class == real_class:\n",
        "      correct +=1\n",
        "    total += 1\n",
        "\n",
        "print(\"Accuracy :\", round(correct/total,3))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "2kOJY94sHDt_"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "LzrUnNHm0gTS"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "WQw0jUAs0mG0"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "mount_file_id": "1OEYBW4RhaTeGKfaskm7JJRAZxm4XdAKP",
      "authorship_tag": "ABX9TyOQ85Sq4CLJMfYNNNBlsGB6",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}