pruned_model_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2a3de92f-6759-41b5-b06b-f832edaec3e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import json\n",
    "import time\n",
    "import statistics\n",
    "import numpy as np\n",
    "import pickle as pkl\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "from transformers.models.bloom.configuration_bloom import BloomConfig\n",
    "from pruning.pruned_bloom import PrunedBloomForCausalLM\n",
    "from node_attribution.utils import count_params\n",
    "from rouge_score import rouge_scorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "42d1fb51-4885-4939-a618-0f96f12b3935",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(f\"bigscience/bloomz-1b1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c0b85f06-12c1-4985-a0c2-548d323b06ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pkl.load(open(\"44_human_filtered_conv_pairs.pkl\", \"rb\"))\n",
    "cali_data = data[:22]\n",
    "val_data = data[22:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "74165a44-7e25-4040-a1a9-6ca5a986be88",
   "metadata": {},
   "outputs": [],
   "source": [
    "def score(model, tokenizer, sentence):\n",
    "    tensor_input = tokenizer.encode(sentence, return_tensors='pt')\n",
    "    repeat_input = tensor_input.repeat(tensor_input.size(-1)-2, 1)\n",
    "    mask = torch.ones(tensor_input.size(-1) - 1).diag(1)[:-2]\n",
    "    masked_input = repeat_input.masked_fill(mask == 1, tokenizer.pad_token_id)\n",
    "    labels = repeat_input.masked_fill( masked_input != tokenizer.pad_token_id, -100)\n",
    "    with torch.inference_mode():\n",
    "        loss = model(masked_input, labels=labels).loss\n",
    "    return np.exp(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8e003e74-29e5-4021-bfd3-f4512a7a4495",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_p(data):\n",
    "    perplexity_sum = 0\n",
    "    for pair in tqdm(data):\n",
    "        perplexity = score(sentence=pair, model=pruned_model, tokenizer=tokenizer)\n",
    "        perplexity_sum += perplexity\n",
    "    \n",
    "    p = perplexity_sum / len(val_data)\n",
    "    \n",
    "    return p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5ed63a74-abb2-4968-aa2c-5a82686d476b",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_path = \"pruned_30percent_560m_bloom.pt\"\n",
    "state_dict_shapes_path = \"pruned_30percent_560m_bloom_state_dict_shapes.pkl\"\n",
    "\n",
    "bloom_config = BloomConfig(\n",
    "    vocab_size=250880,\n",
    "    hidden_size=1024,\n",
    "    n_layer=24,\n",
    "    n_head=16,\n",
    "    layer_norm_epsilon=1e-5,\n",
    "    initializer_range=0.02,\n",
    "    use_cache=True,\n",
    "    bos_token_id=1,\n",
    "    eos_token_id=2,\n",
    "    apply_residual_connection_post_layernorm=False,\n",
    "    hidden_dropout=0.0,\n",
    "    attention_dropout=0.0,\n",
    "    pretraining_tp=1,  # TP rank used when training with megatron\n",
    "    slow_but_exact=False,\n",
    "    attention_softmax_in_fp32=True,\n",
    "    bias_dropout_fusion=True,\n",
    "    masked_softmax_fusion=True,\n",
    "    offset_alibi=100,\n",
    "    pad_token_id=3,\n",
    "    seq_length=2048,\n",
    "    skip_bias_add=True,\n",
    "    skip_bias_add_qkv=False,\n",
    "    unk_token_id=0,\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2d70af46-70b3-4775-b40f-6edccfb70a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "pruned_model = PrunedBloomForCausalLM(bloom_config, state_dict_shapes_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ed6bfc52-6743-48db-82d4-dc8de97a119a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pruned_model.load_state_dict(torch.load(weights_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d4526c93-9061-4e02-9c92-9385086e4294",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.30908961152430015\n"
     ]
    }
   ],
   "source": [
    "pruned_percent = 1.0 - (count_params(pruned_model)[-1] / 559214592)\n",
    "print(pruned_percent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "35f458a6-cc2e-40ab-b522-239a9379252c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|                                                                                                                                                               | 0/22 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([57, 59, 256])\n",
      "1\n",
      "torch.Size([57, 59, 256])\n",
      "2\n",
      "torch.Size([57, 59, 256])\n",
      "3\n",
      "torch.Size([57, 59, 256])\n",
      "4\n",
      "torch.Size([57, 59, 256])\n",
      "5\n",
      "torch.Size([57, 59, 256])\n",
      "6\n",
      "torch.Size([57, 59, 256])\n",
      "7\n",
      "torch.Size([57, 59, 256])\n",
      "8\n",
      "torch.Size([57, 59, 256])\n",
      "9\n",
      "torch.Size([57, 59, 256])\n",
      "10\n",
      "torch.Size([57, 59, 256])\n",
      "11\n",
      "torch.Size([57, 59, 256])\n",
      "12\n",
      "torch.Size([57, 59, 256])\n",
      "13\n",
      "torch.Size([57, 59, 256])\n",
      "14\n",
      "torch.Size([57, 59, 320])\n",
      "15\n",
      "torch.Size([57, 59, 256])\n",
      "16\n",
      "torch.Size([57, 59, 256])\n",
      "17\n",
      "torch.Size([57, 59, 256])\n",
      "18\n",
      "torch.Size([57, 59, 320])\n",
      "19\n",
      "torch.Size([57, 59, 256])\n",
      "20\n",
      "torch.Size([57, 59, 704])\n",
      "21\n",
      "torch.Size([57, 59, 1024])\n",
      "22\n",
      "torch.Size([57, 59, 1024])\n",
      "23\n",
      "torch.Size([57, 59, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  5%|██████▊                                                                                                                                                | 1/22 [00:03<01:09,  3.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([53, 55, 256])\n",
      "1\n",
      "torch.Size([53, 55, 256])\n",
      "2\n",
      "torch.Size([53, 55, 256])\n",
      "3\n",
      "torch.Size([53, 55, 256])\n",
      "4\n",
      "torch.Size([53, 55, 256])\n",
      "5\n",
      "torch.Size([53, 55, 256])\n",
      "6\n",
      "torch.Size([53, 55, 256])\n",
      "7\n",
      "torch.Size([53, 55, 256])\n",
      "8\n",
      "torch.Size([53, 55, 256])\n",
      "9\n",
      "torch.Size([53, 55, 256])\n",
      "10\n",
      "torch.Size([53, 55, 256])\n",
      "11\n",
      "torch.Size([53, 55, 256])\n",
      "12\n",
      "torch.Size([53, 55, 256])\n",
      "13\n",
      "torch.Size([53, 55, 256])\n",
      "14\n",
      "torch.Size([53, 55, 320])\n",
      "15\n",
      "torch.Size([53, 55, 256])\n",
      "16\n",
      "torch.Size([53, 55, 256])\n",
      "17\n",
      "torch.Size([53, 55, 256])\n",
      "18\n",
      "torch.Size([53, 55, 320])\n",
      "19\n",
      "torch.Size([53, 55, 256])\n",
      "20\n",
      "torch.Size([53, 55, 704])\n",
      "21\n",
      "torch.Size([53, 55, 1024])\n",
      "22\n",
      "torch.Size([53, 55, 1024])\n",
      "23\n",
      "torch.Size([53, 55, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  9%|█████████████▋                                                                                                                                         | 2/22 [00:05<00:58,  2.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([64, 66, 256])\n",
      "1\n",
      "torch.Size([64, 66, 256])\n",
      "2\n",
      "torch.Size([64, 66, 256])\n",
      "3\n",
      "torch.Size([64, 66, 256])\n",
      "4\n",
      "torch.Size([64, 66, 256])\n",
      "5\n",
      "torch.Size([64, 66, 256])\n",
      "6\n",
      "torch.Size([64, 66, 256])\n",
      "7\n",
      "torch.Size([64, 66, 256])\n",
      "8\n",
      "torch.Size([64, 66, 256])\n",
      "9\n",
      "torch.Size([64, 66, 256])\n",
      "10\n",
      "torch.Size([64, 66, 256])\n",
      "11\n",
      "torch.Size([64, 66, 256])\n",
      "12\n",
      "torch.Size([64, 66, 256])\n",
      "13\n",
      "torch.Size([64, 66, 256])\n",
      "14\n",
      "torch.Size([64, 66, 320])\n",
      "15\n",
      "torch.Size([64, 66, 256])\n",
      "16\n",
      "torch.Size([64, 66, 256])\n",
      "17\n",
      "torch.Size([64, 66, 256])\n",
      "18\n",
      "torch.Size([64, 66, 320])\n",
      "19\n",
      "torch.Size([64, 66, 256])\n",
      "20\n",
      "torch.Size([64, 66, 704])\n",
      "21\n",
      "torch.Size([64, 66, 1024])\n",
      "22\n",
      "torch.Size([64, 66, 1024])\n",
      "23\n",
      "torch.Size([64, 66, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 14%|████████████████████▌                                                                                                                                  | 3/22 [00:09<01:03,  3.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([79, 81, 256])\n",
      "1\n",
      "torch.Size([79, 81, 256])\n",
      "2\n",
      "torch.Size([79, 81, 256])\n",
      "3\n",
      "torch.Size([79, 81, 256])\n",
      "4\n",
      "torch.Size([79, 81, 256])\n",
      "5\n",
      "torch.Size([79, 81, 256])\n",
      "6\n",
      "torch.Size([79, 81, 256])\n",
      "7\n",
      "torch.Size([79, 81, 256])\n",
      "8\n",
      "torch.Size([79, 81, 256])\n",
      "9\n",
      "torch.Size([79, 81, 256])\n",
      "10\n",
      "torch.Size([79, 81, 256])\n",
      "11\n",
      "torch.Size([79, 81, 256])\n",
      "12\n",
      "torch.Size([79, 81, 256])\n",
      "13\n",
      "torch.Size([79, 81, 256])\n",
      "14\n",
      "torch.Size([79, 81, 320])\n",
      "15\n",
      "torch.Size([79, 81, 256])\n",
      "16\n",
      "torch.Size([79, 81, 256])\n",
      "17\n",
      "torch.Size([79, 81, 256])\n",
      "18\n",
      "torch.Size([79, 81, 320])\n",
      "19\n",
      "torch.Size([79, 81, 256])\n",
      "20\n",
      "torch.Size([79, 81, 704])\n",
      "21\n",
      "torch.Size([79, 81, 1024])\n",
      "22\n",
      "torch.Size([79, 81, 1024])\n",
      "23\n",
      "torch.Size([79, 81, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 18%|███████████████████████████▍                                                                                                                           | 4/22 [00:15<01:20,  4.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([53, 55, 256])\n",
      "1\n",
      "torch.Size([53, 55, 256])\n",
      "2\n",
      "torch.Size([53, 55, 256])\n",
      "3\n",
      "torch.Size([53, 55, 256])\n",
      "4\n",
      "torch.Size([53, 55, 256])\n",
      "5\n",
      "torch.Size([53, 55, 256])\n",
      "6\n",
      "torch.Size([53, 55, 256])\n",
      "7\n",
      "torch.Size([53, 55, 256])\n",
      "8\n",
      "torch.Size([53, 55, 256])\n",
      "9\n",
      "torch.Size([53, 55, 256])\n",
      "10\n",
      "torch.Size([53, 55, 256])\n",
      "11\n",
      "torch.Size([53, 55, 256])\n",
      "12\n",
      "torch.Size([53, 55, 256])\n",
      "13\n",
      "torch.Size([53, 55, 256])\n",
      "14\n",
      "torch.Size([53, 55, 320])\n",
      "15\n",
      "torch.Size([53, 55, 256])\n",
      "16\n",
      "torch.Size([53, 55, 256])\n",
      "17\n",
      "torch.Size([53, 55, 256])\n",
      "18\n",
      "torch.Size([53, 55, 320])\n",
      "19\n",
      "torch.Size([53, 55, 256])\n",
      "20\n",
      "torch.Size([53, 55, 704])\n",
      "21\n",
      "torch.Size([53, 55, 1024])\n",
      "22\n",
      "torch.Size([53, 55, 1024])\n",
      "23\n",
      "torch.Size([53, 55, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 23%|██████████████████████████████████▎                                                                                                                    | 5/22 [00:18<01:05,  3.86s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([104, 106, 256])\n",
      "1\n",
      "torch.Size([104, 106, 256])\n",
      "2\n",
      "torch.Size([104, 106, 256])\n",
      "3\n",
      "torch.Size([104, 106, 256])\n",
      "4\n",
      "torch.Size([104, 106, 256])\n",
      "5\n",
      "torch.Size([104, 106, 256])\n",
      "6\n",
      "torch.Size([104, 106, 256])\n",
      "7\n",
      "torch.Size([104, 106, 256])\n",
      "8\n",
      "torch.Size([104, 106, 256])\n",
      "9\n",
      "torch.Size([104, 106, 256])\n",
      "10\n",
      "torch.Size([104, 106, 256])\n",
      "11\n",
      "torch.Size([104, 106, 256])\n",
      "12\n",
      "torch.Size([104, 106, 256])\n",
      "13\n",
      "torch.Size([104, 106, 256])\n",
      "14\n",
      "torch.Size([104, 106, 320])\n",
      "15\n",
      "torch.Size([104, 106, 256])\n",
      "16\n",
      "torch.Size([104, 106, 256])\n",
      "17\n",
      "torch.Size([104, 106, 256])\n",
      "18\n",
      "torch.Size([104, 106, 320])\n",
      "19\n",
      "torch.Size([104, 106, 256])\n",
      "20\n",
      "torch.Size([104, 106, 704])\n",
      "21\n",
      "torch.Size([104, 106, 1024])\n",
      "22\n",
      "torch.Size([104, 106, 1024])\n",
      "23\n",
      "torch.Size([104, 106, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 27%|█████████████████████████████████████████▏                                                                                                             | 6/22 [00:29<01:39,  6.23s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([60, 62, 256])\n",
      "1\n",
      "torch.Size([60, 62, 256])\n",
      "2\n",
      "torch.Size([60, 62, 256])\n",
      "3\n",
      "torch.Size([60, 62, 256])\n",
      "4\n",
      "torch.Size([60, 62, 256])\n",
      "5\n",
      "torch.Size([60, 62, 256])\n",
      "6\n",
      "torch.Size([60, 62, 256])\n",
      "7\n",
      "torch.Size([60, 62, 256])\n",
      "8\n",
      "torch.Size([60, 62, 256])\n",
      "9\n",
      "torch.Size([60, 62, 256])\n",
      "10\n",
      "torch.Size([60, 62, 256])\n",
      "11\n",
      "torch.Size([60, 62, 256])\n",
      "12\n",
      "torch.Size([60, 62, 256])\n",
      "13\n",
      "torch.Size([60, 62, 256])\n",
      "14\n",
      "torch.Size([60, 62, 320])\n",
      "15\n",
      "torch.Size([60, 62, 256])\n",
      "16\n",
      "torch.Size([60, 62, 256])\n",
      "17\n",
      "torch.Size([60, 62, 256])\n",
      "18\n",
      "torch.Size([60, 62, 320])\n",
      "19\n",
      "torch.Size([60, 62, 256])\n",
      "20\n",
      "torch.Size([60, 62, 704])\n",
      "21\n",
      "torch.Size([60, 62, 1024])\n",
      "22\n",
      "torch.Size([60, 62, 1024])\n",
      "23\n",
      "torch.Size([60, 62, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 32%|████████████████████████████████████████████████                                                                                                       | 7/22 [00:33<01:21,  5.41s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([98, 100, 256])\n",
      "1\n",
      "torch.Size([98, 100, 256])\n",
      "2\n",
      "torch.Size([98, 100, 256])\n",
      "3\n",
      "torch.Size([98, 100, 256])\n",
      "4\n",
      "torch.Size([98, 100, 256])\n",
      "5\n",
      "torch.Size([98, 100, 256])\n",
      "6\n",
      "torch.Size([98, 100, 256])\n",
      "7\n",
      "torch.Size([98, 100, 256])\n",
      "8\n",
      "torch.Size([98, 100, 256])\n",
      "9\n",
      "torch.Size([98, 100, 256])\n",
      "10\n",
      "torch.Size([98, 100, 256])\n",
      "11\n",
      "torch.Size([98, 100, 256])\n",
      "12\n",
      "torch.Size([98, 100, 256])\n",
      "13\n",
      "torch.Size([98, 100, 256])\n",
      "14\n",
      "torch.Size([98, 100, 320])\n",
      "15\n",
      "torch.Size([98, 100, 256])\n",
      "16\n",
      "torch.Size([98, 100, 256])\n",
      "17\n",
      "torch.Size([98, 100, 256])\n",
      "18\n",
      "torch.Size([98, 100, 320])\n",
      "19\n",
      "torch.Size([98, 100, 256])\n",
      "20\n",
      "torch.Size([98, 100, 704])\n",
      "21\n",
      "torch.Size([98, 100, 1024])\n",
      "22\n",
      "torch.Size([98, 100, 1024])\n",
      "23\n",
      "torch.Size([98, 100, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 36%|██████████████████████████████████████████████████████▉                                                                                                | 8/22 [00:42<01:33,  6.67s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([48, 50, 256])\n",
      "1\n",
      "torch.Size([48, 50, 256])\n",
      "2\n",
      "torch.Size([48, 50, 256])\n",
      "3\n",
      "torch.Size([48, 50, 256])\n",
      "4\n",
      "torch.Size([48, 50, 256])\n",
      "5\n",
      "torch.Size([48, 50, 256])\n",
      "6\n",
      "torch.Size([48, 50, 256])\n",
      "7\n",
      "torch.Size([48, 50, 256])\n",
      "8\n",
      "torch.Size([48, 50, 256])\n",
      "9\n",
      "torch.Size([48, 50, 256])\n",
      "10\n",
      "torch.Size([48, 50, 256])\n",
      "11\n",
      "torch.Size([48, 50, 256])\n",
      "12\n",
      "torch.Size([48, 50, 256])\n",
      "13\n",
      "torch.Size([48, 50, 256])\n",
      "14\n",
      "torch.Size([48, 50, 320])\n",
      "15\n",
      "torch.Size([48, 50, 256])\n",
      "16\n",
      "torch.Size([48, 50, 256])\n",
      "17\n",
      "torch.Size([48, 50, 256])\n",
      "18\n",
      "torch.Size([48, 50, 320])\n",
      "19\n",
      "torch.Size([48, 50, 256])\n",
      "20\n",
      "torch.Size([48, 50, 704])\n",
      "21\n",
      "torch.Size([48, 50, 1024])\n",
      "22\n",
      "torch.Size([48, 50, 1024])\n",
      "23\n",
      "torch.Size([48, 50, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 41%|█████████████████████████████████████████████████████████████▊                                                                                         | 9/22 [00:44<01:08,  5.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([81, 83, 256])\n",
      "1\n",
      "torch.Size([81, 83, 256])\n",
      "2\n",
      "torch.Size([81, 83, 256])\n",
      "3\n",
      "torch.Size([81, 83, 256])\n",
      "4\n",
      "torch.Size([81, 83, 256])\n",
      "5\n",
      "torch.Size([81, 83, 256])\n",
      "6\n",
      "torch.Size([81, 83, 256])\n",
      "7\n",
      "torch.Size([81, 83, 256])\n",
      "8\n",
      "torch.Size([81, 83, 256])\n",
      "9\n",
      "torch.Size([81, 83, 256])\n",
      "10\n",
      "torch.Size([81, 83, 256])\n",
      "11\n",
      "torch.Size([81, 83, 256])\n",
      "12\n",
      "torch.Size([81, 83, 256])\n",
      "13\n",
      "torch.Size([81, 83, 256])\n",
      "14\n",
      "torch.Size([81, 83, 320])\n",
      "15\n",
      "torch.Size([81, 83, 256])\n",
      "16\n",
      "torch.Size([81, 83, 256])\n",
      "17\n",
      "torch.Size([81, 83, 256])\n",
      "18\n",
      "torch.Size([81, 83, 320])\n",
      "19\n",
      "torch.Size([81, 83, 256])\n",
      "20\n",
      "torch.Size([81, 83, 704])\n",
      "21\n",
      "torch.Size([81, 83, 1024])\n",
      "22\n",
      "torch.Size([81, 83, 1024])\n",
      "23\n",
      "torch.Size([81, 83, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 45%|████████████████████████████████████████████████████████████████████▏                                                                                 | 10/22 [00:51<01:07,  5.66s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([86, 88, 256])\n",
      "1\n",
      "torch.Size([86, 88, 256])\n",
      "2\n",
      "torch.Size([86, 88, 256])\n",
      "3\n",
      "torch.Size([86, 88, 256])\n",
      "4\n",
      "torch.Size([86, 88, 256])\n",
      "5\n",
      "torch.Size([86, 88, 256])\n",
      "6\n",
      "torch.Size([86, 88, 256])\n",
      "7\n",
      "torch.Size([86, 88, 256])\n",
      "8\n",
      "torch.Size([86, 88, 256])\n",
      "9\n",
      "torch.Size([86, 88, 256])\n",
      "10\n",
      "torch.Size([86, 88, 256])\n",
      "11\n",
      "torch.Size([86, 88, 256])\n",
      "12\n",
      "torch.Size([86, 88, 256])\n",
      "13\n",
      "torch.Size([86, 88, 256])\n",
      "14\n",
      "torch.Size([86, 88, 320])\n",
      "15\n",
      "torch.Size([86, 88, 256])\n",
      "16\n",
      "torch.Size([86, 88, 256])\n",
      "17\n",
      "torch.Size([86, 88, 256])\n",
      "18\n",
      "torch.Size([86, 88, 320])\n",
      "19\n",
      "torch.Size([86, 88, 256])\n",
      "20\n",
      "torch.Size([86, 88, 704])\n",
      "21\n",
      "torch.Size([86, 88, 1024])\n",
      "22\n",
      "torch.Size([86, 88, 1024])\n",
      "23\n",
      "torch.Size([86, 88, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 50%|███████████████████████████████████████████████████████████████████████████                                                                           | 11/22 [00:58<01:07,  6.13s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([66, 68, 256])\n",
      "1\n",
      "torch.Size([66, 68, 256])\n",
      "2\n",
      "torch.Size([66, 68, 256])\n",
      "3\n",
      "torch.Size([66, 68, 256])\n",
      "4\n",
      "torch.Size([66, 68, 256])\n",
      "5\n",
      "torch.Size([66, 68, 256])\n",
      "6\n",
      "torch.Size([66, 68, 256])\n",
      "7\n",
      "torch.Size([66, 68, 256])\n",
      "8\n",
      "torch.Size([66, 68, 256])\n",
      "9\n",
      "torch.Size([66, 68, 256])\n",
      "10\n",
      "torch.Size([66, 68, 256])\n",
      "11\n",
      "torch.Size([66, 68, 256])\n",
      "12\n",
      "torch.Size([66, 68, 256])\n",
      "13\n",
      "torch.Size([66, 68, 256])\n",
      "14\n",
      "torch.Size([66, 68, 320])\n",
      "15\n",
      "torch.Size([66, 68, 256])\n",
      "16\n",
      "torch.Size([66, 68, 256])\n",
      "17\n",
      "torch.Size([66, 68, 256])\n",
      "18\n",
      "torch.Size([66, 68, 320])\n",
      "19\n",
      "torch.Size([66, 68, 256])\n",
      "20\n",
      "torch.Size([66, 68, 704])\n",
      "21\n",
      "torch.Size([66, 68, 1024])\n",
      "22\n",
      "torch.Size([66, 68, 1024])\n",
      "23\n",
      "torch.Size([66, 68, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 55%|█████████████████████████████████████████████████████████████████████████████████▊                                                                    | 12/22 [01:02<00:55,  5.53s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([127, 129, 256])\n",
      "1\n",
      "torch.Size([127, 129, 256])\n",
      "2\n",
      "torch.Size([127, 129, 256])\n",
      "3\n",
      "torch.Size([127, 129, 256])\n",
      "4\n",
      "torch.Size([127, 129, 256])\n",
      "5\n",
      "torch.Size([127, 129, 256])\n",
      "6\n",
      "torch.Size([127, 129, 256])\n",
      "7\n",
      "torch.Size([127, 129, 256])\n",
      "8\n",
      "torch.Size([127, 129, 256])\n",
      "9\n",
      "torch.Size([127, 129, 256])\n",
      "10\n",
      "torch.Size([127, 129, 256])\n",
      "11\n",
      "torch.Size([127, 129, 256])\n",
      "12\n",
      "torch.Size([127, 129, 256])\n",
      "13\n",
      "torch.Size([127, 129, 256])\n",
      "14\n",
      "torch.Size([127, 129, 320])\n",
      "15\n",
      "torch.Size([127, 129, 256])\n",
      "16\n",
      "torch.Size([127, 129, 256])\n",
      "17\n",
      "torch.Size([127, 129, 256])\n",
      "18\n",
      "torch.Size([127, 129, 320])\n",
      "19\n",
      "torch.Size([127, 129, 256])\n",
      "20\n",
      "torch.Size([127, 129, 704])\n",
      "21\n",
      "torch.Size([127, 129, 1024])\n",
      "22\n",
      "torch.Size([127, 129, 1024])\n",
      "23\n",
      "torch.Size([127, 129, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 59%|████████████████████████████████████████████████████████████████████████████████████████▋                                                             | 13/22 [01:34<02:00, 13.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([89, 91, 256])\n",
      "1\n",
      "torch.Size([89, 91, 256])\n",
      "2\n",
      "torch.Size([89, 91, 256])\n",
      "3\n",
      "torch.Size([89, 91, 256])\n",
      "4\n",
      "torch.Size([89, 91, 256])\n",
      "5\n",
      "torch.Size([89, 91, 256])\n",
      "6\n",
      "torch.Size([89, 91, 256])\n",
      "7\n",
      "torch.Size([89, 91, 256])\n",
      "8\n",
      "torch.Size([89, 91, 256])\n",
      "9\n",
      "torch.Size([89, 91, 256])\n",
      "10\n",
      "torch.Size([89, 91, 256])\n",
      "11\n",
      "torch.Size([89, 91, 256])\n",
      "12\n",
      "torch.Size([89, 91, 256])\n",
      "13\n",
      "torch.Size([89, 91, 256])\n",
      "14\n",
      "torch.Size([89, 91, 320])\n",
      "15\n",
      "torch.Size([89, 91, 256])\n",
      "16\n",
      "torch.Size([89, 91, 256])\n",
      "17\n",
      "torch.Size([89, 91, 256])\n",
      "18\n",
      "torch.Size([89, 91, 320])\n",
      "19\n",
      "torch.Size([89, 91, 256])\n",
      "20\n",
      "torch.Size([89, 91, 704])\n",
      "21\n",
      "torch.Size([89, 91, 1024])\n",
      "22\n",
      "torch.Size([89, 91, 1024])\n",
      "23\n",
      "torch.Size([89, 91, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 64%|███████████████████████████████████████████████████████████████████████████████████████████████▍                                                      | 14/22 [01:42<01:33, 11.73s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([126, 128, 256])\n",
      "1\n",
      "torch.Size([126, 128, 256])\n",
      "2\n",
      "torch.Size([126, 128, 256])\n",
      "3\n",
      "torch.Size([126, 128, 256])\n",
      "4\n",
      "torch.Size([126, 128, 256])\n",
      "5\n",
      "torch.Size([126, 128, 256])\n",
      "6\n",
      "torch.Size([126, 128, 256])\n",
      "7\n",
      "torch.Size([126, 128, 256])\n",
      "8\n",
      "torch.Size([126, 128, 256])\n",
      "9\n",
      "torch.Size([126, 128, 256])\n",
      "10\n",
      "torch.Size([126, 128, 256])\n",
      "11\n",
      "torch.Size([126, 128, 256])\n",
      "12\n",
      "torch.Size([126, 128, 256])\n",
      "13\n",
      "torch.Size([126, 128, 256])\n",
      "14\n",
      "torch.Size([126, 128, 320])\n",
      "15\n",
      "torch.Size([126, 128, 256])\n",
      "16\n",
      "torch.Size([126, 128, 256])\n",
      "17\n",
      "torch.Size([126, 128, 256])\n",
      "18\n",
      "torch.Size([126, 128, 320])\n",
      "19\n",
      "torch.Size([126, 128, 256])\n",
      "20\n",
      "torch.Size([126, 128, 704])\n",
      "21\n",
      "torch.Size([126, 128, 1024])\n",
      "22\n",
      "torch.Size([126, 128, 1024])\n",
      "23\n",
      "torch.Size([126, 128, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 68%|██████████████████████████████████████████████████████████████████████████████████████████████████████▎                                               | 15/22 [02:11<01:59, 17.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([117, 119, 256])\n",
      "1\n",
      "torch.Size([117, 119, 256])\n",
      "2\n",
      "torch.Size([117, 119, 256])\n",
      "3\n",
      "torch.Size([117, 119, 256])\n",
      "4\n",
      "torch.Size([117, 119, 256])\n",
      "5\n",
      "torch.Size([117, 119, 256])\n",
      "6\n",
      "torch.Size([117, 119, 256])\n",
      "7\n",
      "torch.Size([117, 119, 256])\n",
      "8\n",
      "torch.Size([117, 119, 256])\n",
      "9\n",
      "torch.Size([117, 119, 256])\n",
      "10\n",
      "torch.Size([117, 119, 256])\n",
      "11\n",
      "torch.Size([117, 119, 256])\n",
      "12\n",
      "torch.Size([117, 119, 256])\n",
      "13\n",
      "torch.Size([117, 119, 256])\n",
      "14\n",
      "torch.Size([117, 119, 320])\n",
      "15\n",
      "torch.Size([117, 119, 256])\n",
      "16\n",
      "torch.Size([117, 119, 256])\n",
      "17\n",
      "torch.Size([117, 119, 256])\n",
      "18\n",
      "torch.Size([117, 119, 320])\n",
      "19\n",
      "torch.Size([117, 119, 256])\n",
      "20\n",
      "torch.Size([117, 119, 704])\n",
      "21\n",
      "torch.Size([117, 119, 1024])\n",
      "22\n",
      "torch.Size([117, 119, 1024])\n",
      "23\n",
      "torch.Size([117, 119, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 73%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████                                         | 16/22 [02:24<01:34, 15.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([120, 122, 256])\n",
      "1\n",
      "torch.Size([120, 122, 256])\n",
      "2\n",
      "torch.Size([120, 122, 256])\n",
      "3\n",
      "torch.Size([120, 122, 256])\n",
      "4\n",
      "torch.Size([120, 122, 256])\n",
      "5\n",
      "torch.Size([120, 122, 256])\n",
      "6\n",
      "torch.Size([120, 122, 256])\n",
      "7\n",
      "torch.Size([120, 122, 256])\n",
      "8\n",
      "torch.Size([120, 122, 256])\n",
      "9\n",
      "torch.Size([120, 122, 256])\n",
      "10\n",
      "torch.Size([120, 122, 256])\n",
      "11\n",
      "torch.Size([120, 122, 256])\n",
      "12\n",
      "torch.Size([120, 122, 256])\n",
      "13\n",
      "torch.Size([120, 122, 256])\n",
      "14\n",
      "torch.Size([120, 122, 320])\n",
      "15\n",
      "torch.Size([120, 122, 256])\n",
      "16\n",
      "torch.Size([120, 122, 256])\n",
      "17\n",
      "torch.Size([120, 122, 256])\n",
      "18\n",
      "torch.Size([120, 122, 320])\n",
      "19\n",
      "torch.Size([120, 122, 256])\n",
      "20\n",
      "torch.Size([120, 122, 704])\n",
      "21\n",
      "torch.Size([120, 122, 1024])\n",
      "22\n",
      "torch.Size([120, 122, 1024])\n",
      "23\n",
      "torch.Size([120, 122, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 77%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉                                  | 17/22 [02:38<01:15, 15.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([59, 61, 256])\n",
      "1\n",
      "torch.Size([59, 61, 256])\n",
      "2\n",
      "torch.Size([59, 61, 256])\n",
      "3\n",
      "torch.Size([59, 61, 256])\n",
      "4\n",
      "torch.Size([59, 61, 256])\n",
      "5\n",
      "torch.Size([59, 61, 256])\n",
      "6\n",
      "torch.Size([59, 61, 256])\n",
      "7\n",
      "torch.Size([59, 61, 256])\n",
      "8\n",
      "torch.Size([59, 61, 256])\n",
      "9\n",
      "torch.Size([59, 61, 256])\n",
      "10\n",
      "torch.Size([59, 61, 256])\n",
      "11\n",
      "torch.Size([59, 61, 256])\n",
      "12\n",
      "torch.Size([59, 61, 256])\n",
      "13\n",
      "torch.Size([59, 61, 256])\n",
      "14\n",
      "torch.Size([59, 61, 320])\n",
      "15\n",
      "torch.Size([59, 61, 256])\n",
      "16\n",
      "torch.Size([59, 61, 256])\n",
      "17\n",
      "torch.Size([59, 61, 256])\n",
      "18\n",
      "torch.Size([59, 61, 320])\n",
      "19\n",
      "torch.Size([59, 61, 256])\n",
      "20\n",
      "torch.Size([59, 61, 704])\n",
      "21\n",
      "torch.Size([59, 61, 1024])\n",
      "22\n",
      "torch.Size([59, 61, 1024])\n",
      "23\n",
      "torch.Size([59, 61, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 82%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                           | 18/22 [02:41<00:46, 11.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([70, 72, 256])\n",
      "1\n",
      "torch.Size([70, 72, 256])\n",
      "2\n",
      "torch.Size([70, 72, 256])\n",
      "3\n",
      "torch.Size([70, 72, 256])\n",
      "4\n",
      "torch.Size([70, 72, 256])\n",
      "5\n",
      "torch.Size([70, 72, 256])\n",
      "6\n",
      "torch.Size([70, 72, 256])\n",
      "7\n",
      "torch.Size([70, 72, 256])\n",
      "8\n",
      "torch.Size([70, 72, 256])\n",
      "9\n",
      "torch.Size([70, 72, 256])\n",
      "10\n",
      "torch.Size([70, 72, 256])\n",
      "11\n",
      "torch.Size([70, 72, 256])\n",
      "12\n",
      "torch.Size([70, 72, 256])\n",
      "13\n",
      "torch.Size([70, 72, 256])\n",
      "14\n",
      "torch.Size([70, 72, 320])\n",
      "15\n",
      "torch.Size([70, 72, 256])\n",
      "16\n",
      "torch.Size([70, 72, 256])\n",
      "17\n",
      "torch.Size([70, 72, 256])\n",
      "18\n",
      "torch.Size([70, 72, 320])\n",
      "19\n",
      "torch.Size([70, 72, 256])\n",
      "20\n",
      "torch.Size([70, 72, 704])\n",
      "21\n",
      "torch.Size([70, 72, 1024])\n",
      "22\n",
      "torch.Size([70, 72, 1024])\n",
      "23\n",
      "torch.Size([70, 72, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 86%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                    | 19/22 [02:46<00:28,  9.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([69, 71, 256])\n",
      "1\n",
      "torch.Size([69, 71, 256])\n",
      "2\n",
      "torch.Size([69, 71, 256])\n",
      "3\n",
      "torch.Size([69, 71, 256])\n",
      "4\n",
      "torch.Size([69, 71, 256])\n",
      "5\n",
      "torch.Size([69, 71, 256])\n",
      "6\n",
      "torch.Size([69, 71, 256])\n",
      "7\n",
      "torch.Size([69, 71, 256])\n",
      "8\n",
      "torch.Size([69, 71, 256])\n",
      "9\n",
      "torch.Size([69, 71, 256])\n",
      "10\n",
      "torch.Size([69, 71, 256])\n",
      "11\n",
      "torch.Size([69, 71, 256])\n",
      "12\n",
      "torch.Size([69, 71, 256])\n",
      "13\n",
      "torch.Size([69, 71, 256])\n",
      "14\n",
      "torch.Size([69, 71, 320])\n",
      "15\n",
      "torch.Size([69, 71, 256])\n",
      "16\n",
      "torch.Size([69, 71, 256])\n",
      "17\n",
      "torch.Size([69, 71, 256])\n",
      "18\n",
      "torch.Size([69, 71, 320])\n",
      "19\n",
      "torch.Size([69, 71, 256])\n",
      "20\n",
      "torch.Size([69, 71, 704])\n",
      "21\n",
      "torch.Size([69, 71, 1024])\n",
      "22\n",
      "torch.Size([69, 71, 1024])\n",
      "23\n",
      "torch.Size([69, 71, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 91%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎             | 20/22 [02:50<00:15,  7.95s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([104, 106, 256])\n",
      "1\n",
      "torch.Size([104, 106, 256])\n",
      "2\n",
      "torch.Size([104, 106, 256])\n",
      "3\n",
      "torch.Size([104, 106, 256])\n",
      "4\n",
      "torch.Size([104, 106, 256])\n",
      "5\n",
      "torch.Size([104, 106, 256])\n",
      "6\n",
      "torch.Size([104, 106, 256])\n",
      "7\n",
      "torch.Size([104, 106, 256])\n",
      "8\n",
      "torch.Size([104, 106, 256])\n",
      "9\n",
      "torch.Size([104, 106, 256])\n",
      "10\n",
      "torch.Size([104, 106, 256])\n",
      "11\n",
      "torch.Size([104, 106, 256])\n",
      "12\n",
      "torch.Size([104, 106, 256])\n",
      "13\n",
      "torch.Size([104, 106, 256])\n",
      "14\n",
      "torch.Size([104, 106, 320])\n",
      "15\n",
      "torch.Size([104, 106, 256])\n",
      "16\n",
      "torch.Size([104, 106, 256])\n",
      "17\n",
      "torch.Size([104, 106, 256])\n",
      "18\n",
      "torch.Size([104, 106, 320])\n",
      "19\n",
      "torch.Size([104, 106, 256])\n",
      "20\n",
      "torch.Size([104, 106, 704])\n",
      "21\n",
      "torch.Size([104, 106, 1024])\n",
      "22\n",
      "torch.Size([104, 106, 1024])\n",
      "23\n",
      "torch.Size([104, 106, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 95%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏      | 21/22 [03:01<00:08,  8.74s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([91, 93, 256])\n",
      "1\n",
      "torch.Size([91, 93, 256])\n",
      "2\n",
      "torch.Size([91, 93, 256])\n",
      "3\n",
      "torch.Size([91, 93, 256])\n",
      "4\n",
      "torch.Size([91, 93, 256])\n",
      "5\n",
      "torch.Size([91, 93, 256])\n",
      "6\n",
      "torch.Size([91, 93, 256])\n",
      "7\n",
      "torch.Size([91, 93, 256])\n",
      "8\n",
      "torch.Size([91, 93, 256])\n",
      "9\n",
      "torch.Size([91, 93, 256])\n",
      "10\n",
      "torch.Size([91, 93, 256])\n",
      "11\n",
      "torch.Size([91, 93, 256])\n",
      "12\n",
      "torch.Size([91, 93, 256])\n",
      "13\n",
      "torch.Size([91, 93, 256])\n",
      "14\n",
      "torch.Size([91, 93, 320])\n",
      "15\n",
      "torch.Size([91, 93, 256])\n",
      "16\n",
      "torch.Size([91, 93, 256])\n",
      "17\n",
      "torch.Size([91, 93, 256])\n",
      "18\n",
      "torch.Size([91, 93, 320])\n",
      "19\n",
      "torch.Size([91, 93, 256])\n",
      "20\n",
      "torch.Size([91, 93, 704])\n",
      "21\n",
      "torch.Size([91, 93, 1024])\n",
      "22\n",
      "torch.Size([91, 93, 1024])\n",
      "23\n",
      "torch.Size([91, 93, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [03:09<00:00,  8.60s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21432584.78894007\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|                                                                                                                                                               | 0/22 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([86, 88, 256])\n",
      "1\n",
      "torch.Size([86, 88, 256])\n",
      "2\n",
      "torch.Size([86, 88, 256])\n",
      "3\n",
      "torch.Size([86, 88, 256])\n",
      "4\n",
      "torch.Size([86, 88, 256])\n",
      "5\n",
      "torch.Size([86, 88, 256])\n",
      "6\n",
      "torch.Size([86, 88, 256])\n",
      "7\n",
      "torch.Size([86, 88, 256])\n",
      "8\n",
      "torch.Size([86, 88, 256])\n",
      "9\n",
      "torch.Size([86, 88, 256])\n",
      "10\n",
      "torch.Size([86, 88, 256])\n",
      "11\n",
      "torch.Size([86, 88, 256])\n",
      "12\n",
      "torch.Size([86, 88, 256])\n",
      "13\n",
      "torch.Size([86, 88, 256])\n",
      "14\n",
      "torch.Size([86, 88, 320])\n",
      "15\n",
      "torch.Size([86, 88, 256])\n",
      "16\n",
      "torch.Size([86, 88, 256])\n",
      "17\n",
      "torch.Size([86, 88, 256])\n",
      "18\n",
      "torch.Size([86, 88, 320])\n",
      "19\n",
      "torch.Size([86, 88, 256])\n",
      "20\n",
      "torch.Size([86, 88, 704])\n",
      "21\n",
      "torch.Size([86, 88, 1024])\n",
      "22\n",
      "torch.Size([86, 88, 1024])\n",
      "23\n",
      "torch.Size([86, 88, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  5%|██████▊                                                                                                                                                | 1/22 [00:07<02:30,  7.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([98, 100, 256])\n",
      "1\n",
      "torch.Size([98, 100, 256])\n",
      "2\n",
      "torch.Size([98, 100, 256])\n",
      "3\n",
      "torch.Size([98, 100, 256])\n",
      "4\n",
      "torch.Size([98, 100, 256])\n",
      "5\n",
      "torch.Size([98, 100, 256])\n",
      "6\n",
      "torch.Size([98, 100, 256])\n",
      "7\n",
      "torch.Size([98, 100, 256])\n",
      "8\n",
      "torch.Size([98, 100, 256])\n",
      "9\n",
      "torch.Size([98, 100, 256])\n",
      "10\n",
      "torch.Size([98, 100, 256])\n",
      "11\n",
      "torch.Size([98, 100, 256])\n",
      "12\n",
      "torch.Size([98, 100, 256])\n",
      "13\n",
      "torch.Size([98, 100, 256])\n",
      "14\n",
      "torch.Size([98, 100, 320])\n",
      "15\n",
      "torch.Size([98, 100, 256])\n",
      "16\n",
      "torch.Size([98, 100, 256])\n",
      "17\n",
      "torch.Size([98, 100, 256])\n",
      "18\n",
      "torch.Size([98, 100, 320])\n",
      "19\n",
      "torch.Size([98, 100, 256])\n",
      "20\n",
      "torch.Size([98, 100, 704])\n",
      "21\n",
      "torch.Size([98, 100, 1024])\n",
      "22\n",
      "torch.Size([98, 100, 1024])\n",
      "23\n",
      "torch.Size([98, 100, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  9%|█████████████▋                                                                                                                                         | 2/22 [00:16<02:46,  8.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([103, 105, 256])\n",
      "1\n",
      "torch.Size([103, 105, 256])\n",
      "2\n",
      "torch.Size([103, 105, 256])\n",
      "3\n",
      "torch.Size([103, 105, 256])\n",
      "4\n",
      "torch.Size([103, 105, 256])\n",
      "5\n",
      "torch.Size([103, 105, 256])\n",
      "6\n",
      "torch.Size([103, 105, 256])\n",
      "7\n",
      "torch.Size([103, 105, 256])\n",
      "8\n",
      "torch.Size([103, 105, 256])\n",
      "9\n",
      "torch.Size([103, 105, 256])\n",
      "10\n",
      "torch.Size([103, 105, 256])\n",
      "11\n",
      "torch.Size([103, 105, 256])\n",
      "12\n",
      "torch.Size([103, 105, 256])\n",
      "13\n",
      "torch.Size([103, 105, 256])\n",
      "14\n",
      "torch.Size([103, 105, 320])\n",
      "15\n",
      "torch.Size([103, 105, 256])\n",
      "16\n",
      "torch.Size([103, 105, 256])\n",
      "17\n",
      "torch.Size([103, 105, 256])\n",
      "18\n",
      "torch.Size([103, 105, 320])\n",
      "19\n",
      "torch.Size([103, 105, 256])\n",
      "20\n",
      "torch.Size([103, 105, 704])\n",
      "21\n",
      "torch.Size([103, 105, 1024])\n",
      "22\n",
      "torch.Size([103, 105, 1024])\n",
      "23\n",
      "torch.Size([103, 105, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 14%|████████████████████▌                                                                                                                                  | 3/22 [00:26<02:54,  9.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([67, 69, 256])\n",
      "1\n",
      "torch.Size([67, 69, 256])\n",
      "2\n",
      "torch.Size([67, 69, 256])\n",
      "3\n",
      "torch.Size([67, 69, 256])\n",
      "4\n",
      "torch.Size([67, 69, 256])\n",
      "5\n",
      "torch.Size([67, 69, 256])\n",
      "6\n",
      "torch.Size([67, 69, 256])\n",
      "7\n",
      "torch.Size([67, 69, 256])\n",
      "8\n",
      "torch.Size([67, 69, 256])\n",
      "9\n",
      "torch.Size([67, 69, 256])\n",
      "10\n",
      "torch.Size([67, 69, 256])\n",
      "11\n",
      "torch.Size([67, 69, 256])\n",
      "12\n",
      "torch.Size([67, 69, 256])\n",
      "13\n",
      "torch.Size([67, 69, 256])\n",
      "14\n",
      "torch.Size([67, 69, 320])\n",
      "15\n",
      "torch.Size([67, 69, 256])\n",
      "16\n",
      "torch.Size([67, 69, 256])\n",
      "17\n",
      "torch.Size([67, 69, 256])\n",
      "18\n",
      "torch.Size([67, 69, 320])\n",
      "19\n",
      "torch.Size([67, 69, 256])\n",
      "20\n",
      "torch.Size([67, 69, 704])\n",
      "21\n",
      "torch.Size([67, 69, 1024])\n",
      "22\n",
      "torch.Size([67, 69, 1024])\n",
      "23\n",
      "torch.Size([67, 69, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 18%|███████████████████████████▍                                                                                                                           | 4/22 [00:30<02:11,  7.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([62, 64, 256])\n",
      "1\n",
      "torch.Size([62, 64, 256])\n",
      "2\n",
      "torch.Size([62, 64, 256])\n",
      "3\n",
      "torch.Size([62, 64, 256])\n",
      "4\n",
      "torch.Size([62, 64, 256])\n",
      "5\n",
      "torch.Size([62, 64, 256])\n",
      "6\n",
      "torch.Size([62, 64, 256])\n",
      "7\n",
      "torch.Size([62, 64, 256])\n",
      "8\n",
      "torch.Size([62, 64, 256])\n",
      "9\n",
      "torch.Size([62, 64, 256])\n",
      "10\n",
      "torch.Size([62, 64, 256])\n",
      "11\n",
      "torch.Size([62, 64, 256])\n",
      "12\n",
      "torch.Size([62, 64, 256])\n",
      "13\n",
      "torch.Size([62, 64, 256])\n",
      "14\n",
      "torch.Size([62, 64, 320])\n",
      "15\n",
      "torch.Size([62, 64, 256])\n",
      "16\n",
      "torch.Size([62, 64, 256])\n",
      "17\n",
      "torch.Size([62, 64, 256])\n",
      "18\n",
      "torch.Size([62, 64, 320])\n",
      "19\n",
      "torch.Size([62, 64, 256])\n",
      "20\n",
      "torch.Size([62, 64, 704])\n",
      "21\n",
      "torch.Size([62, 64, 1024])\n",
      "22\n",
      "torch.Size([62, 64, 1024])\n",
      "23\n",
      "torch.Size([62, 64, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 23%|██████████████████████████████████▎                                                                                                                    | 5/22 [00:34<01:42,  6.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([39, 41, 256])\n",
      "1\n",
      "torch.Size([39, 41, 256])\n",
      "2\n",
      "torch.Size([39, 41, 256])\n",
      "3\n",
      "torch.Size([39, 41, 256])\n",
      "4\n",
      "torch.Size([39, 41, 256])\n",
      "5\n",
      "torch.Size([39, 41, 256])\n",
      "6\n",
      "torch.Size([39, 41, 256])\n",
      "7\n",
      "torch.Size([39, 41, 256])\n",
      "8\n",
      "torch.Size([39, 41, 256])\n",
      "9\n",
      "torch.Size([39, 41, 256])\n",
      "10\n",
      "torch.Size([39, 41, 256])\n",
      "11\n",
      "torch.Size([39, 41, 256])\n",
      "12\n",
      "torch.Size([39, 41, 256])\n",
      "13\n",
      "torch.Size([39, 41, 256])\n",
      "14\n",
      "torch.Size([39, 41, 320])\n",
      "15\n",
      "torch.Size([39, 41, 256])\n",
      "16\n",
      "torch.Size([39, 41, 256])\n",
      "17\n",
      "torch.Size([39, 41, 256])\n",
      "18\n",
      "torch.Size([39, 41, 320])\n",
      "19\n",
      "torch.Size([39, 41, 256])\n",
      "20\n",
      "torch.Size([39, 41, 704])\n",
      "21\n",
      "torch.Size([39, 41, 1024])\n",
      "22\n",
      "torch.Size([39, 41, 1024])\n",
      "23\n",
      "torch.Size([39, 41, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 27%|█████████████████████████████████████████▏                                                                                                             | 6/22 [00:36<01:13,  4.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([39, 41, 256])\n",
      "1\n",
      "torch.Size([39, 41, 256])\n",
      "2\n",
      "torch.Size([39, 41, 256])\n",
      "3\n",
      "torch.Size([39, 41, 256])\n",
      "4\n",
      "torch.Size([39, 41, 256])\n",
      "5\n",
      "torch.Size([39, 41, 256])\n",
      "6\n",
      "torch.Size([39, 41, 256])\n",
      "7\n",
      "torch.Size([39, 41, 256])\n",
      "8\n",
      "torch.Size([39, 41, 256])\n",
      "9\n",
      "torch.Size([39, 41, 256])\n",
      "10\n",
      "torch.Size([39, 41, 256])\n",
      "11\n",
      "torch.Size([39, 41, 256])\n",
      "12\n",
      "torch.Size([39, 41, 256])\n",
      "13\n",
      "torch.Size([39, 41, 256])\n",
      "14\n",
      "torch.Size([39, 41, 320])\n",
      "15\n",
      "torch.Size([39, 41, 256])\n",
      "16\n",
      "torch.Size([39, 41, 256])\n",
      "17\n",
      "torch.Size([39, 41, 256])\n",
      "18\n",
      "torch.Size([39, 41, 320])\n",
      "19\n",
      "torch.Size([39, 41, 256])\n",
      "20\n",
      "torch.Size([39, 41, 704])\n",
      "21\n",
      "torch.Size([39, 41, 1024])\n",
      "22\n",
      "torch.Size([39, 41, 1024])\n",
      "23\n",
      "torch.Size([39, 41, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 32%|████████████████████████████████████████████████                                                                                                       | 7/22 [00:38<00:54,  3.60s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([68, 70, 256])\n",
      "1\n",
      "torch.Size([68, 70, 256])\n",
      "2\n",
      "torch.Size([68, 70, 256])\n",
      "3\n",
      "torch.Size([68, 70, 256])\n",
      "4\n",
      "torch.Size([68, 70, 256])\n",
      "5\n",
      "torch.Size([68, 70, 256])\n",
      "6\n",
      "torch.Size([68, 70, 256])\n",
      "7\n",
      "torch.Size([68, 70, 256])\n",
      "8\n",
      "torch.Size([68, 70, 256])\n",
      "9\n",
      "torch.Size([68, 70, 256])\n",
      "10\n",
      "torch.Size([68, 70, 256])\n",
      "11\n",
      "torch.Size([68, 70, 256])\n",
      "12\n",
      "torch.Size([68, 70, 256])\n",
      "13\n",
      "torch.Size([68, 70, 256])\n",
      "14\n",
      "torch.Size([68, 70, 320])\n",
      "15\n",
      "torch.Size([68, 70, 256])\n",
      "16\n",
      "torch.Size([68, 70, 256])\n",
      "17\n",
      "torch.Size([68, 70, 256])\n",
      "18\n",
      "torch.Size([68, 70, 320])\n",
      "19\n",
      "torch.Size([68, 70, 256])\n",
      "20\n",
      "torch.Size([68, 70, 704])\n",
      "21\n",
      "torch.Size([68, 70, 1024])\n",
      "22\n",
      "torch.Size([68, 70, 1024])\n",
      "23\n",
      "torch.Size([68, 70, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 36%|██████████████████████████████████████████████████████▉                                                                                                | 8/22 [00:42<00:54,  3.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([75, 77, 256])\n",
      "1\n",
      "torch.Size([75, 77, 256])\n",
      "2\n",
      "torch.Size([75, 77, 256])\n",
      "3\n",
      "torch.Size([75, 77, 256])\n",
      "4\n",
      "torch.Size([75, 77, 256])\n",
      "5\n",
      "torch.Size([75, 77, 256])\n",
      "6\n",
      "torch.Size([75, 77, 256])\n",
      "7\n",
      "torch.Size([75, 77, 256])\n",
      "8\n",
      "torch.Size([75, 77, 256])\n",
      "9\n",
      "torch.Size([75, 77, 256])\n",
      "10\n",
      "torch.Size([75, 77, 256])\n",
      "11\n",
      "torch.Size([75, 77, 256])\n",
      "12\n",
      "torch.Size([75, 77, 256])\n",
      "13\n",
      "torch.Size([75, 77, 256])\n",
      "14\n",
      "torch.Size([75, 77, 320])\n",
      "15\n",
      "torch.Size([75, 77, 256])\n",
      "16\n",
      "torch.Size([75, 77, 256])\n",
      "17\n",
      "torch.Size([75, 77, 256])\n",
      "18\n",
      "torch.Size([75, 77, 320])\n",
      "19\n",
      "torch.Size([75, 77, 256])\n",
      "20\n",
      "torch.Size([75, 77, 704])\n",
      "21\n",
      "torch.Size([75, 77, 1024])\n",
      "22\n",
      "torch.Size([75, 77, 1024])\n",
      "23\n",
      "torch.Size([75, 77, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 41%|█████████████████████████████████████████████████████████████▊                                                                                         | 9/22 [00:47<00:56,  4.37s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([79, 81, 256])\n",
      "1\n",
      "torch.Size([79, 81, 256])\n",
      "2\n",
      "torch.Size([79, 81, 256])\n",
      "3\n",
      "torch.Size([79, 81, 256])\n",
      "4\n",
      "torch.Size([79, 81, 256])\n",
      "5\n",
      "torch.Size([79, 81, 256])\n",
      "6\n",
      "torch.Size([79, 81, 256])\n",
      "7\n",
      "torch.Size([79, 81, 256])\n",
      "8\n",
      "torch.Size([79, 81, 256])\n",
      "9\n",
      "torch.Size([79, 81, 256])\n",
      "10\n",
      "torch.Size([79, 81, 256])\n",
      "11\n",
      "torch.Size([79, 81, 256])\n",
      "12\n",
      "torch.Size([79, 81, 256])\n",
      "13\n",
      "torch.Size([79, 81, 256])\n",
      "14\n",
      "torch.Size([79, 81, 320])\n",
      "15\n",
      "torch.Size([79, 81, 256])\n",
      "16\n",
      "torch.Size([79, 81, 256])\n",
      "17\n",
      "torch.Size([79, 81, 256])\n",
      "18\n",
      "torch.Size([79, 81, 320])\n",
      "19\n",
      "torch.Size([79, 81, 256])\n",
      "20\n",
      "torch.Size([79, 81, 704])\n",
      "21\n",
      "torch.Size([79, 81, 1024])\n",
      "22\n",
      "torch.Size([79, 81, 1024])\n",
      "23\n",
      "torch.Size([79, 81, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 45%|████████████████████████████████████████████████████████████████████▏                                                                                 | 10/22 [00:54<00:58,  4.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([64, 66, 256])\n",
      "1\n",
      "torch.Size([64, 66, 256])\n",
      "2\n",
      "torch.Size([64, 66, 256])\n",
      "3\n",
      "torch.Size([64, 66, 256])\n",
      "4\n",
      "torch.Size([64, 66, 256])\n",
      "5\n",
      "torch.Size([64, 66, 256])\n",
      "6\n",
      "torch.Size([64, 66, 256])\n",
      "7\n",
      "torch.Size([64, 66, 256])\n",
      "8\n",
      "torch.Size([64, 66, 256])\n",
      "9\n",
      "torch.Size([64, 66, 256])\n",
      "10\n",
      "torch.Size([64, 66, 256])\n",
      "11\n",
      "torch.Size([64, 66, 256])\n",
      "12\n",
      "torch.Size([64, 66, 256])\n",
      "13\n",
      "torch.Size([64, 66, 256])\n",
      "14\n",
      "torch.Size([64, 66, 320])\n",
      "15\n",
      "torch.Size([64, 66, 256])\n",
      "16\n",
      "torch.Size([64, 66, 256])\n",
      "17\n",
      "torch.Size([64, 66, 256])\n",
      "18\n",
      "torch.Size([64, 66, 320])\n",
      "19\n",
      "torch.Size([64, 66, 256])\n",
      "20\n",
      "torch.Size([64, 66, 704])\n",
      "21\n",
      "torch.Size([64, 66, 1024])\n",
      "22\n",
      "torch.Size([64, 66, 1024])\n",
      "23\n",
      "torch.Size([64, 66, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 50%|███████████████████████████████████████████████████████████████████████████                                                                           | 11/22 [00:57<00:50,  4.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([53, 55, 256])\n",
      "1\n",
      "torch.Size([53, 55, 256])\n",
      "2\n",
      "torch.Size([53, 55, 256])\n",
      "3\n",
      "torch.Size([53, 55, 256])\n",
      "4\n",
      "torch.Size([53, 55, 256])\n",
      "5\n",
      "torch.Size([53, 55, 256])\n",
      "6\n",
      "torch.Size([53, 55, 256])\n",
      "7\n",
      "torch.Size([53, 55, 256])\n",
      "8\n",
      "torch.Size([53, 55, 256])\n",
      "9\n",
      "torch.Size([53, 55, 256])\n",
      "10\n",
      "torch.Size([53, 55, 256])\n",
      "11\n",
      "torch.Size([53, 55, 256])\n",
      "12\n",
      "torch.Size([53, 55, 256])\n",
      "13\n",
      "torch.Size([53, 55, 256])\n",
      "14\n",
      "torch.Size([53, 55, 320])\n",
      "15\n",
      "torch.Size([53, 55, 256])\n",
      "16\n",
      "torch.Size([53, 55, 256])\n",
      "17\n",
      "torch.Size([53, 55, 256])\n",
      "18\n",
      "torch.Size([53, 55, 320])\n",
      "19\n",
      "torch.Size([53, 55, 256])\n",
      "20\n",
      "torch.Size([53, 55, 704])\n",
      "21\n",
      "torch.Size([53, 55, 1024])\n",
      "22\n",
      "torch.Size([53, 55, 1024])\n",
      "23\n",
      "torch.Size([53, 55, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 55%|█████████████████████████████████████████████████████████████████████████████████▊                                                                    | 12/22 [01:00<00:40,  4.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([52, 54, 256])\n",
      "1\n",
      "torch.Size([52, 54, 256])\n",
      "2\n",
      "torch.Size([52, 54, 256])\n",
      "3\n",
      "torch.Size([52, 54, 256])\n",
      "4\n",
      "torch.Size([52, 54, 256])\n",
      "5\n",
      "torch.Size([52, 54, 256])\n",
      "6\n",
      "torch.Size([52, 54, 256])\n",
      "7\n",
      "torch.Size([52, 54, 256])\n",
      "8\n",
      "torch.Size([52, 54, 256])\n",
      "9\n",
      "torch.Size([52, 54, 256])\n",
      "10\n",
      "torch.Size([52, 54, 256])\n",
      "11\n",
      "torch.Size([52, 54, 256])\n",
      "12\n",
      "torch.Size([52, 54, 256])\n",
      "13\n",
      "torch.Size([52, 54, 256])\n",
      "14\n",
      "torch.Size([52, 54, 320])\n",
      "15\n",
      "torch.Size([52, 54, 256])\n",
      "16\n",
      "torch.Size([52, 54, 256])\n",
      "17\n",
      "torch.Size([52, 54, 256])\n",
      "18\n",
      "torch.Size([52, 54, 320])\n",
      "19\n",
      "torch.Size([52, 54, 256])\n",
      "20\n",
      "torch.Size([52, 54, 704])\n",
      "21\n",
      "torch.Size([52, 54, 1024])\n",
      "22\n",
      "torch.Size([52, 54, 1024])\n",
      "23\n",
      "torch.Size([52, 54, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 59%|████████████████████████████████████████████████████████████████████████████████████████▋                                                             | 13/22 [01:03<00:32,  3.58s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([91, 93, 256])\n",
      "1\n",
      "torch.Size([91, 93, 256])\n",
      "2\n",
      "torch.Size([91, 93, 256])\n",
      "3\n",
      "torch.Size([91, 93, 256])\n",
      "4\n",
      "torch.Size([91, 93, 256])\n",
      "5\n",
      "torch.Size([91, 93, 256])\n",
      "6\n",
      "torch.Size([91, 93, 256])\n",
      "7\n",
      "torch.Size([91, 93, 256])\n",
      "8\n",
      "torch.Size([91, 93, 256])\n",
      "9\n",
      "torch.Size([91, 93, 256])\n",
      "10\n",
      "torch.Size([91, 93, 256])\n",
      "11\n",
      "torch.Size([91, 93, 256])\n",
      "12\n",
      "torch.Size([91, 93, 256])\n",
      "13\n",
      "torch.Size([91, 93, 256])\n",
      "14\n",
      "torch.Size([91, 93, 320])\n",
      "15\n",
      "torch.Size([91, 93, 256])\n",
      "16\n",
      "torch.Size([91, 93, 256])\n",
      "17\n",
      "torch.Size([91, 93, 256])\n",
      "18\n",
      "torch.Size([91, 93, 320])\n",
      "19\n",
      "torch.Size([91, 93, 256])\n",
      "20\n",
      "torch.Size([91, 93, 704])\n",
      "21\n",
      "torch.Size([91, 93, 1024])\n",
      "22\n",
      "torch.Size([91, 93, 1024])\n",
      "23\n",
      "torch.Size([91, 93, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 64%|███████████████████████████████████████████████████████████████████████████████████████████████▍                                                      | 14/22 [01:11<00:39,  4.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([40, 42, 256])\n",
      "1\n",
      "torch.Size([40, 42, 256])\n",
      "2\n",
      "torch.Size([40, 42, 256])\n",
      "3\n",
      "torch.Size([40, 42, 256])\n",
      "4\n",
      "torch.Size([40, 42, 256])\n",
      "5\n",
      "torch.Size([40, 42, 256])\n",
      "6\n",
      "torch.Size([40, 42, 256])\n",
      "7\n",
      "torch.Size([40, 42, 256])\n",
      "8\n",
      "torch.Size([40, 42, 256])\n",
      "9\n",
      "torch.Size([40, 42, 256])\n",
      "10\n",
      "torch.Size([40, 42, 256])\n",
      "11\n",
      "torch.Size([40, 42, 256])\n",
      "12\n",
      "torch.Size([40, 42, 256])\n",
      "13\n",
      "torch.Size([40, 42, 256])\n",
      "14\n",
      "torch.Size([40, 42, 320])\n",
      "15\n",
      "torch.Size([40, 42, 256])\n",
      "16\n",
      "torch.Size([40, 42, 256])\n",
      "17\n",
      "torch.Size([40, 42, 256])\n",
      "18\n",
      "torch.Size([40, 42, 320])\n",
      "19\n",
      "torch.Size([40, 42, 256])\n",
      "20\n",
      "torch.Size([40, 42, 704])\n",
      "21\n",
      "torch.Size([40, 42, 1024])\n",
      "22\n",
      "torch.Size([40, 42, 1024])\n",
      "23\n",
      "torch.Size([40, 42, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 68%|██████████████████████████████████████████████████████████████████████████████████████████████████████▎                                               | 15/22 [01:12<00:27,  3.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([46, 48, 256])\n",
      "1\n",
      "torch.Size([46, 48, 256])\n",
      "2\n",
      "torch.Size([46, 48, 256])\n",
      "3\n",
      "torch.Size([46, 48, 256])\n",
      "4\n",
      "torch.Size([46, 48, 256])\n",
      "5\n",
      "torch.Size([46, 48, 256])\n",
      "6\n",
      "torch.Size([46, 48, 256])\n",
      "7\n",
      "torch.Size([46, 48, 256])\n",
      "8\n",
      "torch.Size([46, 48, 256])\n",
      "9\n",
      "torch.Size([46, 48, 256])\n",
      "10\n",
      "torch.Size([46, 48, 256])\n",
      "11\n",
      "torch.Size([46, 48, 256])\n",
      "12\n",
      "torch.Size([46, 48, 256])\n",
      "13\n",
      "torch.Size([46, 48, 256])\n",
      "14\n",
      "torch.Size([46, 48, 320])\n",
      "15\n",
      "torch.Size([46, 48, 256])\n",
      "16\n",
      "torch.Size([46, 48, 256])\n",
      "17\n",
      "torch.Size([46, 48, 256])\n",
      "18\n",
      "torch.Size([46, 48, 320])\n",
      "19\n",
      "torch.Size([46, 48, 256])\n",
      "20\n",
      "torch.Size([46, 48, 704])\n",
      "21\n",
      "torch.Size([46, 48, 1024])\n",
      "22\n",
      "torch.Size([46, 48, 1024])\n",
      "23\n",
      "torch.Size([46, 48, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 73%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████                                         | 16/22 [01:14<00:20,  3.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([84, 86, 256])\n",
      "1\n",
      "torch.Size([84, 86, 256])\n",
      "2\n",
      "torch.Size([84, 86, 256])\n",
      "3\n",
      "torch.Size([84, 86, 256])\n",
      "4\n",
      "torch.Size([84, 86, 256])\n",
      "5\n",
      "torch.Size([84, 86, 256])\n",
      "6\n",
      "torch.Size([84, 86, 256])\n",
      "7\n",
      "torch.Size([84, 86, 256])\n",
      "8\n",
      "torch.Size([84, 86, 256])\n",
      "9\n",
      "torch.Size([84, 86, 256])\n",
      "10\n",
      "torch.Size([84, 86, 256])\n",
      "11\n",
      "torch.Size([84, 86, 256])\n",
      "12\n",
      "torch.Size([84, 86, 256])\n",
      "13\n",
      "torch.Size([84, 86, 256])\n",
      "14\n",
      "torch.Size([84, 86, 320])\n",
      "15\n",
      "torch.Size([84, 86, 256])\n",
      "16\n",
      "torch.Size([84, 86, 256])\n",
      "17\n",
      "torch.Size([84, 86, 256])\n",
      "18\n",
      "torch.Size([84, 86, 320])\n",
      "19\n",
      "torch.Size([84, 86, 256])\n",
      "20\n",
      "torch.Size([84, 86, 704])\n",
      "21\n",
      "torch.Size([84, 86, 1024])\n",
      "22\n",
      "torch.Size([84, 86, 1024])\n",
      "23\n",
      "torch.Size([84, 86, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 77%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉                                  | 17/22 [01:21<00:21,  4.34s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([77, 79, 256])\n",
      "1\n",
      "torch.Size([77, 79, 256])\n",
      "2\n",
      "torch.Size([77, 79, 256])\n",
      "3\n",
      "torch.Size([77, 79, 256])\n",
      "4\n",
      "torch.Size([77, 79, 256])\n",
      "5\n",
      "torch.Size([77, 79, 256])\n",
      "6\n",
      "torch.Size([77, 79, 256])\n",
      "7\n",
      "torch.Size([77, 79, 256])\n",
      "8\n",
      "torch.Size([77, 79, 256])\n",
      "9\n",
      "torch.Size([77, 79, 256])\n",
      "10\n",
      "torch.Size([77, 79, 256])\n",
      "11\n",
      "torch.Size([77, 79, 256])\n",
      "12\n",
      "torch.Size([77, 79, 256])\n",
      "13\n",
      "torch.Size([77, 79, 256])\n",
      "14\n",
      "torch.Size([77, 79, 320])\n",
      "15\n",
      "torch.Size([77, 79, 256])\n",
      "16\n",
      "torch.Size([77, 79, 256])\n",
      "17\n",
      "torch.Size([77, 79, 256])\n",
      "18\n",
      "torch.Size([77, 79, 320])\n",
      "19\n",
      "torch.Size([77, 79, 256])\n",
      "20\n",
      "torch.Size([77, 79, 704])\n",
      "21\n",
      "torch.Size([77, 79, 1024])\n",
      "22\n",
      "torch.Size([77, 79, 1024])\n",
      "23\n",
      "torch.Size([77, 79, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 82%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                           | 18/22 [01:26<00:18,  4.68s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([97, 99, 256])\n",
      "1\n",
      "torch.Size([97, 99, 256])\n",
      "2\n",
      "torch.Size([97, 99, 256])\n",
      "3\n",
      "torch.Size([97, 99, 256])\n",
      "4\n",
      "torch.Size([97, 99, 256])\n",
      "5\n",
      "torch.Size([97, 99, 256])\n",
      "6\n",
      "torch.Size([97, 99, 256])\n",
      "7\n",
      "torch.Size([97, 99, 256])\n",
      "8\n",
      "torch.Size([97, 99, 256])\n",
      "9\n",
      "torch.Size([97, 99, 256])\n",
      "10\n",
      "torch.Size([97, 99, 256])\n",
      "11\n",
      "torch.Size([97, 99, 256])\n",
      "12\n",
      "torch.Size([97, 99, 256])\n",
      "13\n",
      "torch.Size([97, 99, 256])\n",
      "14\n",
      "torch.Size([97, 99, 320])\n",
      "15\n",
      "torch.Size([97, 99, 256])\n",
      "16\n",
      "torch.Size([97, 99, 256])\n",
      "17\n",
      "torch.Size([97, 99, 256])\n",
      "18\n",
      "torch.Size([97, 99, 320])\n",
      "19\n",
      "torch.Size([97, 99, 256])\n",
      "20\n",
      "torch.Size([97, 99, 704])\n",
      "21\n",
      "torch.Size([97, 99, 1024])\n",
      "22\n",
      "torch.Size([97, 99, 1024])\n",
      "23\n",
      "torch.Size([97, 99, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 86%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                    | 19/22 [01:35<00:17,  5.94s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([96, 98, 256])\n",
      "1\n",
      "torch.Size([96, 98, 256])\n",
      "2\n",
      "torch.Size([96, 98, 256])\n",
      "3\n",
      "torch.Size([96, 98, 256])\n",
      "4\n",
      "torch.Size([96, 98, 256])\n",
      "5\n",
      "torch.Size([96, 98, 256])\n",
      "6\n",
      "torch.Size([96, 98, 256])\n",
      "7\n",
      "torch.Size([96, 98, 256])\n",
      "8\n",
      "torch.Size([96, 98, 256])\n",
      "9\n",
      "torch.Size([96, 98, 256])\n",
      "10\n",
      "torch.Size([96, 98, 256])\n",
      "11\n",
      "torch.Size([96, 98, 256])\n",
      "12\n",
      "torch.Size([96, 98, 256])\n",
      "13\n",
      "torch.Size([96, 98, 256])\n",
      "14\n",
      "torch.Size([96, 98, 320])\n",
      "15\n",
      "torch.Size([96, 98, 256])\n",
      "16\n",
      "torch.Size([96, 98, 256])\n",
      "17\n",
      "torch.Size([96, 98, 256])\n",
      "18\n",
      "torch.Size([96, 98, 320])\n",
      "19\n",
      "torch.Size([96, 98, 256])\n",
      "20\n",
      "torch.Size([96, 98, 704])\n",
      "21\n",
      "torch.Size([96, 98, 1024])\n",
      "22\n",
      "torch.Size([96, 98, 1024])\n",
      "23\n",
      "torch.Size([96, 98, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 91%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎             | 20/22 [01:44<00:13,  6.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([123, 125, 256])\n",
      "1\n",
      "torch.Size([123, 125, 256])\n",
      "2\n",
      "torch.Size([123, 125, 256])\n",
      "3\n",
      "torch.Size([123, 125, 256])\n",
      "4\n",
      "torch.Size([123, 125, 256])\n",
      "5\n",
      "torch.Size([123, 125, 256])\n",
      "6\n",
      "torch.Size([123, 125, 256])\n",
      "7\n",
      "torch.Size([123, 125, 256])\n",
      "8\n",
      "torch.Size([123, 125, 256])\n",
      "9\n",
      "torch.Size([123, 125, 256])\n",
      "10\n",
      "torch.Size([123, 125, 256])\n",
      "11\n",
      "torch.Size([123, 125, 256])\n",
      "12\n",
      "torch.Size([123, 125, 256])\n",
      "13\n",
      "torch.Size([123, 125, 256])\n",
      "14\n",
      "torch.Size([123, 125, 320])\n",
      "15\n",
      "torch.Size([123, 125, 256])\n",
      "16\n",
      "torch.Size([123, 125, 256])\n",
      "17\n",
      "torch.Size([123, 125, 256])\n",
      "18\n",
      "torch.Size([123, 125, 320])\n",
      "19\n",
      "torch.Size([123, 125, 256])\n",
      "20\n",
      "torch.Size([123, 125, 704])\n",
      "21\n",
      "torch.Size([123, 125, 1024])\n",
      "22\n",
      "torch.Size([123, 125, 1024])\n",
      "23\n",
      "torch.Size([123, 125, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 95%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏      | 21/22 [02:11<00:12, 12.72s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "torch.Size([116, 118, 256])\n",
      "1\n",
      "torch.Size([116, 118, 256])\n",
      "2\n",
      "torch.Size([116, 118, 256])\n",
      "3\n",
      "torch.Size([116, 118, 256])\n",
      "4\n",
      "torch.Size([116, 118, 256])\n",
      "5\n",
      "torch.Size([116, 118, 256])\n",
      "6\n",
      "torch.Size([116, 118, 256])\n",
      "7\n",
      "torch.Size([116, 118, 256])\n",
      "8\n",
      "torch.Size([116, 118, 256])\n",
      "9\n",
      "torch.Size([116, 118, 256])\n",
      "10\n",
      "torch.Size([116, 118, 256])\n",
      "11\n",
      "torch.Size([116, 118, 256])\n",
      "12\n",
      "torch.Size([116, 118, 256])\n",
      "13\n",
      "torch.Size([116, 118, 256])\n",
      "14\n",
      "torch.Size([116, 118, 320])\n",
      "15\n",
      "torch.Size([116, 118, 256])\n",
      "16\n",
      "torch.Size([116, 118, 256])\n",
      "17\n",
      "torch.Size([116, 118, 256])\n",
      "18\n",
      "torch.Size([116, 118, 320])\n",
      "19\n",
      "torch.Size([116, 118, 256])\n",
      "20\n",
      "torch.Size([116, 118, 704])\n",
      "21\n",
      "torch.Size([116, 118, 1024])\n",
      "22\n",
      "torch.Size([116, 118, 1024])\n",
      "23\n",
      "torch.Size([116, 118, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [02:23<00:00,  6.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27434652.018668972\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "print(calc_p(cali_data))\n",
    "print(calc_p(val_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0db0af9e-3a66-4d72-88eb-f7b638582f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trials = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0446e0b-5008-4cf0-8e2d-23e6ab6ad16b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "line = \"Hello, I am a social bot! \"\n",
    "inputs = tokenizer(line, return_tensors=\"pt\")\n",
    "pruned_times = []\n",
    "\n",
    "for i in range(num_trials):\n",
    "    start = time.time()\n",
    "    outputs = pruned_model.generate(\n",
    "        input_ids=inputs[\"input_ids\"], \n",
    "        max_new_tokens=100, \n",
    "        do_sample=True, \n",
    "        top_k=50, \n",
    "        top_p=0.95,\n",
    "    )\n",
    "    end = time.time()\n",
    "    pruned_times.append(end - start)\n",
    "    print(f\"inference time: {end - start}\")\n",
    "    \n",
    "print(tokenizer.batch_decode(outputs, skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73f2b340-809f-488b-9f6c-e6fca8f0fdd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_pruned = np.mean(pruned_times)\n",
    "std_pruned = np.std(pruned_times)\n",
    "mode = statistics.median(pruned_times)\n",
    "print(mean_pruned)\n",
    "print(std_pruned)\n",
    "print(mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "104a4c81-3a20-4efd-bd97-c79efc2e31ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# while True:\n",
    "#     line = input(\"You:\")\n",
    "#     inputs = tokenizer(line, return_tensors=\"pt\")\n",
    "#     outputs = pruned_model.generate(\n",
    "#         input_ids=inputs[\"input_ids\"], \n",
    "#         max_new_tokens=20, \n",
    "#         do_sample=True, \n",
    "#         top_k=50, \n",
    "#         top_p=0.95,\n",
    "#     )\n",
    "#     print(tokenizer.batch_decode(outputs, skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "155b46d0-210c-4021-a9f8-02d332687cd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# full_line = \"Person: My favorite movie is the The Day After Tomorrow\\nSocialBot: Oh, interesting, I am not familiar with that movie! Can you tell me more about it?\"\n",
    "# prompt_line = \"Person: My favorite movie is The Day After Tomorrow\\nSocialBot: \"\n",
    "# completion = full_line.split(prompt_line)[-1]\n",
    "# inputs = tokenizer(prompt_line, return_tensors=\"pt\")\n",
    "\n",
    "# #for i in range(num_trials):\n",
    "# outputs = pruned_model.generate(\n",
    "#     input_ids=inputs[\"input_ids\"], \n",
    "#     max_new_tokens=25, \n",
    "#     do_sample=True, \n",
    "#     top_k=50, \n",
    "#     top_p=0.95,\n",
    "# )\n",
    "# out_seq = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "# out_seq = out_seq[0].split(\"Person: My favorite movie is The Day After Tomorrow\\nSocialBot: \")[-1]\n",
    "# r_scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)\n",
    "# rouge_scores = r_scorer.score(completion, out_seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a587c8-e2d2-4204-8daf-3068a718fa17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from transformers import AutoTokenizer, BloomForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4bb5cd7-6d04-4876-ac9f-5dda689bd22e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = BloomForCausalLM.from_pretrained(\"bigscience/bloom-560m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78c663a4-c791-45e2-a9d1-53f79aaa3013",
   "metadata": {},
   "outputs": [],
   "source": [
    "# inputs = tokenizer(\"Person: My favorite movie is The Day After Tomorrow\\nSocialBot: \", return_tensors=\"pt\")\n",
    "# start = time.time()\n",
    "# outputs = model.generate(input_ids=inputs[\"input_ids\"], max_new_tokens=25, do_sample=True, top_k=50, top_p=0.95)\n",
    "# end = time.time()\n",
    "# print(tokenizer.batch_decode(outputs, skip_special_tokens=True))\n",
    "# print(f\"inference time: {end - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e57c06d6-9aa1-467e-a8bc-809884a49555",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tap",
   "language": "python",
   "name": "tap"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}