diff --git a/examples/sparse_mnist.ipynb b/examples/sparse_mnist.ipynb
index 11ccb3f..4c8fa20 100644
--- a/examples/sparse_mnist.ipynb
+++ b/examples/sparse_mnist.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,8 @@
     "# with Numenta, Inc., for a separate license for this software code, the\n",
     "# following terms and conditions apply:\n",
     "#\n",
-    "# This program is free software: you can redistribute it and/or modify\n",
+    "# This program is free software: \n",
+    "# you can redistribute it and/or modify\n",
     "# it under the terms of the GNU Affero Public License version 3 as\n",
     "# published by the Free Software Foundation.\n",
     "#\n",
@@ -30,9 +31,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 13,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Collecting nupic.torch\n",
+      "  Cloning https://github.com/numenta/nupic.torch.git to /private/var/folders/0k/nx51st811lj2bkn32z3c0_dc0000gn/T/pip-install-h0b3r322/nupic-torch_6389283cd1294fb79c8a7c3c4e32176d\n",
+      "  Running command git clone --filter=blob:none --quiet https://github.com/numenta/nupic.torch.git /private/var/folders/0k/nx51st811lj2bkn32z3c0_dc0000gn/T/pip-install-h0b3r322/nupic-torch_6389283cd1294fb79c8a7c3c4e32176d\n",
+      "  Resolved https://github.com/numenta/nupic.torch.git to commit 51269b7aae1024ee301837320a3b3ebee538c8c0\n",
+      "  Installing build dependencies ... \u001b[?25ldone\n",
+      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
+      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
+      "\u001b[?25hRequirement already satisfied: torch<=1.11.0,>=1.6 in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from nupic.torch) (1.11.0)\n",
+      "Requirement already satisfied: typing-extensions in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from torch<=1.11.0,>=1.6->nupic.torch) (4.3.0)\n",
+      "Requirement already satisfied: torch in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (1.11.0)\n",
+      "Requirement already satisfied: torchvision in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (0.13.0)\n",
+      "Requirement already satisfied: typing-extensions in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from torch) (4.3.0)\n",
+      "Requirement already satisfied: numpy in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from torchvision) (1.23.1)\n",
+      "Requirement already satisfied: requests in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from torchvision) (2.28.1)\n",
+      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from torchvision) (9.2.0)\n",
+      "Requirement already satisfied: charset-normalizer<3,>=2 in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from requests->torchvision) (2.1.0)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from requests->torchvision) (2022.6.15)\n",
+      "Requirement already satisfied: idna<4,>=2.5 in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from requests->torchvision) (3.3)\n",
+      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /Users/bard/miniforge3/envs/nupic.torch/lib/python3.8/site-packages (from requests->torchvision) (1.26.11)\n"
+     ]
+    }
+   ],
    "source": [
     "!pip install git+https://github.com/numenta/nupic.torch.git#egg=nupic.torch\n",
     "!pip install torch torchvision"
@@ -40,7 +67,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -51,7 +78,8 @@
     "import torch.nn.functional as F\n",
     "import torch.optim as optim\n",
     "from torchvision import datasets, transforms\n",
-    "from tqdm import tqdm_notebook as tqdm\n",
+    "# from tqdm import tqdm_notebook as tqdm\n",
+    "from tqdm.notebook import tqdm\n",
     "\n",
     "SEED = 18\n",
     "random.seed(SEED)\n",
@@ -64,7 +92,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -132,7 +160,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -159,7 +187,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
@@ -172,21 +200,21 @@
       "    (module): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1))\n",
       "  )\n",
       "  (cnn1_maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
-      "  (cnn1_kwinner): KWinners2d(channels=32, local=False, n=0, percent_on=0.1, boost_strength=1.5, boost_strength_factor=0.85, k_inference_factor=1.0, duty_cycle_period=1000)\n",
+      "  (cnn1_kwinner): KWinners2d(channels=32, local=False, break_ties=False, n=0, percent_on=0.1, boost_strength=1.5, boost_strength_factor=0.85, k_inference_factor=1.0, duty_cycle_period=1000)\n",
       "  (cnn2): SparseWeights2d(\n",
       "    sparsity=0.55\n",
       "    (module): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))\n",
       "  )\n",
       "  (cnn2_maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
-      "  (cnn2_kwinner): KWinners2d(channels=64, local=False, n=0, percent_on=0.2, boost_strength=1.5, boost_strength_factor=0.85, k_inference_factor=1.0, duty_cycle_period=1000)\n",
+      "  (cnn2_kwinner): KWinners2d(channels=64, local=False, break_ties=False, n=0, percent_on=0.2, boost_strength=1.5, boost_strength_factor=0.85, k_inference_factor=1.0, duty_cycle_period=1000)\n",
       "  (flatten): Flatten()\n",
       "  (linear): SparseWeights(\n",
       "    sparsity=0.8\n",
       "    (module): Linear(in_features=1024, out_features=700, bias=True)\n",
       "  )\n",
-      "  (linear_kwinner): KWinners(n=700, percent_on=0.2, boost_strength=1.5, boost_strength_factor=0.85, k_inference_factor=1.0, duty_cycle_period=1000)\n",
+      "  (linear_kwinner): KWinners(n=700, percent_on=0.2, boost_strength=1.5, boost_strength_factor=0.85, k_inference_factor=1.0, duty_cycle_period=1000, break_ties=False)\n",
       "  (output): Linear(in_features=700, out_features=10, bias=True)\n",
-      "  (softmax): LogSoftmax()\n",
+      "  (softmax): LogSoftmax(dim=1)\n",
       ")\n"
      ]
     }
@@ -208,7 +236,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -232,15 +260,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\r"
-     ]
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/15000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
    "source": [
@@ -264,7 +299,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -281,23 +316,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\r"
-     ]
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
     },
     {
      "data": {
       "text/plain": [
-       "{'accuracy': 0.9823, 'loss': 0.05553661346435547, 'total_correct': 9823}"
+       "{'accuracy': 0.9774, 'loss': 0.06903478813171386, 'total_correct': 9774}"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 21,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -315,27 +357,497 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9906, 'loss': 0.03069741744995117, 'total_correct': 9906}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9907, 'loss': 0.030093585205078124, 'total_correct': 9907}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9916, 'loss': 0.028680802726745606, 'total_correct': 9916}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9914, 'loss': 0.028174154233932495, 'total_correct': 9914}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9908, 'loss': 0.02861804084777832, 'total_correct': 9908}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9914, 'loss': 0.02783290138244629, 'total_correct': 9914}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'accuracy': 0.99, 'loss': 0.03087631359100342, 'total_correct': 9900}\n",
-      "{'accuracy': 0.9913, 'loss': 0.028131033515930177, 'total_correct': 9913}\n",
-      "{'accuracy': 0.9912, 'loss': 0.02829796772003174, 'total_correct': 9912}\n",
-      "{'accuracy': 0.9901, 'loss': 0.03012564697265625, 'total_correct': 9901}\n",
-      "{'accuracy': 0.9909, 'loss': 0.0279187762260437, 'total_correct': 9909}\n",
-      "{'accuracy': 0.9913, 'loss': 0.026879340744018553, 'total_correct': 9913}\n",
-      "{'accuracy': 0.991, 'loss': 0.027533163070678712, 'total_correct': 9910}\n",
-      "{'accuracy': 0.9915, 'loss': 0.026872401237487794, 'total_correct': 9915}\n",
-      "{'accuracy': 0.9916, 'loss': 0.02652479610443115, 'total_correct': 9916}\n",
-      "{'accuracy': 0.9915, 'loss': 0.02673978271484375, 'total_correct': 9915}\n",
-      "{'accuracy': 0.9916, 'loss': 0.026827363300323485, 'total_correct': 9916}\n",
-      "{'accuracy': 0.9914, 'loss': 0.02654646396636963, 'total_correct': 9914}\n",
-      "{'accuracy': 0.9916, 'loss': 0.026357748794555665, 'total_correct': 9916}\n",
-      "{'accuracy': 0.9917, 'loss': 0.026316686820983887, 'total_correct': 9917}\n"
+      "{'accuracy': 0.9914, 'loss': 0.02766239242553711, 'total_correct': 9914}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9914, 'loss': 0.027536948108673097, 'total_correct': 9914}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9917, 'loss': 0.02727954349517822, 'total_correct': 9917}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9915, 'loss': 0.02753580093383789, 'total_correct': 9915}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9917, 'loss': 0.027163607788085936, 'total_correct': 9917}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9921, 'loss': 0.027231089973449706, 'total_correct': 9921}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9921, 'loss': 0.02723951950073242, 'total_correct': 9921}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/938 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'accuracy': 0.9915, 'loss': 0.027270044326782227, 'total_correct': 9915}\n"
      ]
     }
    ],
@@ -359,7 +871,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -386,24 +898,238 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.0 : {'accuracy': 0.9915, 'loss': 0.027270044326782227, 'total_correct': 9915}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.05 : {'accuracy': 0.9897, 'loss': 0.032641384887695316, 'total_correct': 9897}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.1 : {'accuracy': 0.987, 'loss': 0.036816733169555665, 'total_correct': 9870}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.15 : {'accuracy': 0.9851, 'loss': 0.04622526092529297, 'total_correct': 9851}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.2 : {'accuracy': 0.9794, 'loss': 0.06127367515563965, 'total_correct': 9794}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.25 : {'accuracy': 0.9764, 'loss': 0.07657530822753907, 'total_correct': 9764}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.3 : {'accuracy': 0.9669, 'loss': 0.10190569915771484, 'total_correct': 9669}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.35 : {'accuracy': 0.9505, 'loss': 0.15377749786376954, 'total_correct': 9505}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.4 : {'accuracy': 0.9365, 'loss': 0.20168538818359374, 'total_correct': 9365}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.45 : {'accuracy': 0.9088, 'loss': 0.2825202178955078, 'total_correct': 9088}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "0.0 : {'accuracy': 0.9917, 'loss': 0.02631667709350586, 'total_correct': 9917}\n",
-      "0.05 : {'accuracy': 0.9899, 'loss': 0.032843858909606935, 'total_correct': 9899}\n",
-      "0.1 : {'accuracy': 0.9855, 'loss': 0.0439798023223877, 'total_correct': 9855}\n",
-      "0.15 : {'accuracy': 0.9829, 'loss': 0.05501823501586914, 'total_correct': 9829}\n",
-      "0.2 : {'accuracy': 0.976, 'loss': 0.07691587409973144, 'total_correct': 9760}\n",
-      "0.25 : {'accuracy': 0.9655, 'loss': 0.10933124389648438, 'total_correct': 9655}\n",
-      "0.3 : {'accuracy': 0.9499, 'loss': 0.15266261291503908, 'total_correct': 9499}\n",
-      "0.35 : {'accuracy': 0.9305, 'loss': 0.21989019317626954, 'total_correct': 9305}\n",
-      "0.4 : {'accuracy': 0.9015, 'loss': 0.30402026062011717, 'total_correct': 9015}\n",
-      "0.45 : {'accuracy': 0.8543, 'loss': 0.4642761535644531, 'total_correct': 8543}\n",
-      "0.5 : {'accuracy': 0.7951, 'loss': 0.6449000061035156, 'total_correct': 7951}\n"
+      "0.5 : {'accuracy': 0.8725, 'loss': 0.4023683654785156, 'total_correct': 8725}\n"
      ]
     }
    ],
@@ -442,9 +1168,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "nupic.torch",
    "language": "python",
-   "name": "python3"
+   "name": "nupic.torch"
   },
   "language_info": {
    "codemirror_mode": {
@@ -456,9 +1182,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.7"
+   "version": "3.8.13"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }