Repository: jeshraghian/snntorch (language: python)

Issue:
Feature spikegen
-reorganized modules
-latency & rate coding
-spiking CNNs
-Tutorials 1-3


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/python-package-tag.yml b/.github/workflows/python-package-tag.yml
new file mode 100644
index 00000000..cad2dd52
--- /dev/null
+++ b/.github/workflows/python-package-tag.yml
@@ -0,0 +1,44 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: Python package tag
+
+on:
+  push:
+    tags:
+      - '*'
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.6, 3.7, 3.8]
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install tox tox-gh-actions
+    - name: Test with tox
+      run: tox
+
+    #   run: |
+    #     python -m pip install --upgrade pip
+    #     python -m pip install flake8 pytest
+    #     if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    # - name: Lint with flake8
+    #   run: |
+    #     # stop the build if there are Python syntax errors or undefined names
+    #     flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+    #     # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+    #     flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    # - name: Test with pytest
+    #   run: |
+    #     pytest
\ No newline at end of file
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
new file mode 100644
index 00000000..4aa0c8dc
--- /dev/null
+++ b/.github/workflows/python-package.yml
@@ -0,0 +1,45 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: Python package
+
+on:
+  push:
+    branches: [ master, git-actions ]  # remove git-actions once functioning
+  pull_request:
+    branches: [ master, git-actions ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.6, 3.7, 3.8]
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install tox tox-gh-actions
+    - name: Test with tox
+      run: tox
+
+    #   run: |
+    #     python -m pip install --upgrade pip
+    #     python -m pip install flake8 pytest
+    #     if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    # - name: Lint with flake8
+    #   run: |
+    #     # stop the build if there are Python syntax errors or undefined names
+    #     flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+    #     # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+    #     flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    # - name: Test with pytest
+    #   run: |
+    #     pytest
\ No newline at end of file
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
new file mode 100644
index 00000000..7de6a540
--- /dev/null
+++ b/.github/workflows/python-publish.yml
@@ -0,0 +1,31 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: Upload Python Package
+
+on:
+  release: # need this to be tags only on master. Also, only want it to run if python-package is successful.
+    types: [created]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+      run: |
+        python setup.py sdist bdist_wheel
+        twine upload dist/*
\ No newline at end of file
diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index 2fa8feec..00000000
--- a/.travis.yml
+++ /dev/null
@@ -1,17 +0,0 @@
-language: python
-python:
-- 3.8
-- 3.7
-- 3.6
-install: pip install -U tox-travis
-script: tox
-deploy:
-  provider: pypi
-  distributions: sdist bdist_wheel
-  user: __token__
-  password:
-    secure: agCUKFnt1a/UIk+EtM+DV3EkLeQQ4bsDKKMQkqbFnOPOKpKiSJIUgv+iHuw5GHeWkvdFEOpZ8JcSuRbqvTjYm0+DndTB+u6PmzDlgPYMF9BdLTlBBsSQ7b6U/L7Kgtb9wa6EOsSoqGWVRljBP9yN8H6wRV/GOhA/zcvbRdEU7/MHdc2VSHo00EfxhXkphTkSfEjD+TOD2r8mU/I59zRoc+YaaafFoHfZzRKhMNqBOi61AWPp+oyGQ5iRMg+kJxLOtFFJpZhgWhk++fVQ0EddipxbBEXOLhPSW3CZSOmEeXZfCvp97qauarGaEemKkSbRVP0OBlBqriOHov/a29pgm4kbd7zm5gd4O40GUuMEt9/txKuPplIZbspyrtEjKy6rZHHTCoIpwle4ZtHE6/ZK0sn2dIdjUgDoP/cpKki75sAczGJnCOHu8CiAEzBl6XoOmIFz4FVCchHWZnkRfAiygtTo+yw46b9dztXTsZYXo2m43a2KW18qP4LwrPaEKlkgUBs7qKUIiuqb8Ej2oRzz7sx1nwjQ8XvIYyQKgFVS9T80kENtyqQ8dbOGySBHNTTZR0FaudY06wGvvs781CRr8qBcZyUZ/E7Ie4wIB4cUoX5Wy8imfAPHAJGy8g0XAokWB90JiUoj94WfFNKheZg6NsLzhIlYVXZ3LgCWUSXfDqI=
-  on:
-    tags: true
-    repo: jeshraghian/snntorch
-    python: 3.8
diff --git a/README.rst b/README.rst
index 6d585ffb..f624816a 100644
--- a/README.rst
+++ b/README.rst
@@ -6,8 +6,11 @@ Introduction
 .. image:: https://img.shields.io/pypi/v/snntorch.svg
         :target: https://pypi.python.org/pypi/snntorch
 
-.. image:: https://img.shields.io/travis/jeshraghian/snntorch.svg
-        :target: https://travis-ci.com/jeshraghian/snntorch
+.. .. image:: https://img.shields.io/travis/jeshraghian/snntorch.svg
+..         :target: https://travis-ci.com/jeshraghian/snntorch
+
+.. image:: https://github.com/jeshraghian/snntorch/actions/workflows/<WORKFLOW_FILE>/badge.svg
+
 
 .. image:: https://readthedocs.org/projects/snntorch/badge/?version=latest
         :target: https://snntorch.readthedocs.io/en/latest/?badge=latest
@@ -72,18 +75,21 @@ Installation
 
 Run the following to install:
 
-```
+``
 python
+
 pip install snntorch
-```
+``
 
 To install snnTorch from source instead:
 
-```
+``
 git clone https://github.com/jeshraghian/snnTorch
+
 cd snnTorch
+
 python setup.py install
-```
+``
 
 API & Examples 
 ^^^^^^^^^^^^^^^^^^^^^^^^
@@ -114,10 +120,13 @@ If you're ready to contribute to snnTorch, instructions to do so can be `found h
 
 Acknowledgments
 ^^^^^^^^^^^^^^^^^^^^^^^^
-snnTorch was developed by Jason K. Eshraghian in the Lu Group (University of Michigan), with additional contributions from Xinxin Wang and Vincent Sun.
+snnTorch was initially developed by `Jason K. Eshraghian`_ in the Lu Group (University of Michigan).
+
+Additional contributions were made by Xinxin Wang, Vincent Sun, and Emre Neftci.
+
 Several features in snnTorch were inspired by the work of Friedemann Zenke, Emre Neftci, Doo Seok Jeong, Sumit Bam Shrestha and Garrick Orchard.
 
-This work was supported by...
+.. _Jason K. Eshraghian: https://jasoneshraghian.com
 
 License & Copyright
 ^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/_version.py b/_version.py
index 16f47ab3..97e403b3 100644
--- a/_version.py
+++ b/_version.py
@@ -1,2 +1,2 @@
 # fmt: off
-__version__ = '0.1.8'
+__version__ = '0.1.15'
diff --git a/docs/_static/img/examples/visio/equations.docx b/docs/_static/img/examples/visio/equations.docx
index e0770ddd..3e2350f9 100644
Binary files a/docs/_static/img/examples/visio/equations.docx and b/docs/_static/img/examples/visio/equations.docx differ
diff --git a/docs/_static/img/examples/visio/tutorial1.vsdx b/docs/_static/img/examples/visio/tutorial1.vsdx
index 3e751ac1..43aa41fc 100644
Binary files a/docs/_static/img/examples/visio/tutorial1.vsdx and b/docs/_static/img/examples/visio/tutorial1.vsdx differ
diff --git a/docs/_static/img/examples/visio/~$$tutorial1.~vsdx b/docs/_static/img/examples/visio/~$$tutorial1.~vsdx
deleted file mode 100644
index a013d353..00000000
Binary files a/docs/_static/img/examples/visio/~$$tutorial1.~vsdx and /dev/null differ
diff --git a/docs/_static/img/examples/visio/~$uations.docx b/docs/_static/img/examples/visio/~$uations.docx
deleted file mode 100644
index cc1fce96..00000000
Binary files a/docs/_static/img/examples/visio/~$uations.docx and /dev/null differ
diff --git a/docs/conf.py b/docs/conf.py
index 28d3a6b6..9e5a0e3b 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -7,7 +7,7 @@
 import snntorch
 
 # fmt: off
-__version__ = '0.1.8'
+__version__ = '0.1.15'
 # fmt: on
 
 
diff --git a/docs/tutorials.rst b/docs/tutorials.rst
index 838abf85..23b46cf6 100644
--- a/docs/tutorials.rst
+++ b/docs/tutorials.rst
@@ -2,4 +2,37 @@
 Tutorials
 ================
 
-Under construction.
\ No newline at end of file
+The tutorial consists of a series of Google Colab notebooks. Static non-editable versions are also available. 
+
+
+.. list-table::
+   :widths: 20 60 30
+   :header-rows: 1
+
+   * - Tutorial
+     - Title
+     - Colab Link
+   * - `Tutorial 1 <https://colab.research.google.com/github/jeshraghian/snntorch/blob/tutorials/examples/tutorial_1_spikegen.ipynb>`_
+     - Spike Generation and Visualization with spikegen and spikeplot
+     - .. image:: https://colab.research.google.com/assets/colab-badge.svg
+        :alt: Open In Colab
+        :target: https://colab.research.google.com/github/jeshraghian/snntorch/blob/tutorials/examples/tutorial_1_spikegen.ipynb
+
+   * - Tutorial 2
+     - **Under construction:** Neuronal Dynamics with snntorch
+     - .. image:: https://colab.research.google.com/assets/colab-badge.svg
+        :alt: Open In Colab
+
+   * - Tutorial 3
+     - **Under construction:** Deep Learning on static and spiking MNIST
+     - .. image:: https://colab.research.google.com/assets/colab-badge.svg
+        :alt: Open In Colab
+
+
+   * - Tutorial 4
+     - **Under construction:** Convolutional SNNs 
+     - .. image:: https://colab.research.google.com/assets/colab-badge.svg
+        :alt: Open In Colab
+
+
+Future tutorials on spiking neurons and training are scheduled to be completed by April 2021. 
\ No newline at end of file
diff --git a/docs/tutorials/tutorial_1.rst b/docs/tutorials/tutorial_1.rst
new file mode 100644
index 00000000..5dfd2fa8
--- /dev/null
+++ b/docs/tutorials/tutorial_1.rst
@@ -0,0 +1,1 @@
+TO-DO: convert colab notebook to static non-editable rst file.
\ No newline at end of file
diff --git a/examples/tutorial_1_spikegen.ipynb b/examples/tutorial_1_spikegen.ipynb
index 26ac3855..b524ec2f 100644
--- a/examples/tutorial_1_spikegen.ipynb
+++ b/examples/tutorial_1_spikegen.ipynb
@@ -48,7 +48,11 @@
         "\n",
         "# snnTorch - Spike Generation and Visualization with `spikegen` and `spikeplot`\n",
         "## Tutorial 1\n",
-        "### By Jason K. Eshraghian"
+        "### By Jason K. Eshraghian\n",
+        "\n",
+        "<a href=\"https://colab.research.google.com/github/jeshraghian/snntorch/blob/tutorials/examples/tutorial_1_spikegen.ipynb\">\n",
+        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+        "</a>"
       ]
     },
     {
diff --git a/examples/tutorial_2_FCN.ipynb b/examples/tutorial_2_FCN.ipynb
index 451ec8ff..6ce4ad96 100644
--- a/examples/tutorial_2_FCN.ipynb
+++ b/examples/tutorial_2_FCN.ipynb
@@ -1,1019 +1,1055 @@
 {
- "cells": [
-  {
-   "cell_type": "markdown",
-   "source": [
-    "# snnTorch - Tutorial 2\n",
-    "### By Jason K. Eshraghian"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "kernelspec": {
+      "name": "python3",
+      "language": "python",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 2
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython2",
+      "version": "2.7.6"
+    },
+    "colab": {
+      "name": "tutorial_2_FCN.ipynb",
+      "provenance": []
     }
-   }
   },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Gradient-based Learning in Spiking Neural Networks\n",
-    "In this tutorial, we'll create a simple 2-layer fully-connected network (FCN) to classify the MNIST dataset.\n",
-    "We will use the backpropagation through time (BPTT) algorithm to train the network.\n",
-    "\n",
-    "If running in Google Colab:\n",
-    "* Ensure you are connected to GPU by checking Runtime > Change runtime type > Hardware accelerator: GPU\n",
-    "* Next, install the Test PyPi distribution of snnTorch by clicking into the following cell and pressing `Shift+Enter`."
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# Install the test PyPi Distribution of snntorch\n",
-    "!pip install -i https://test.pypi.org/simple/ snntorch"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 1. Setting up the Static MNIST Dataset\n",
-    "### 1.1. Import packages and setup environment"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "import snntorch as snn\n",
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "from torch.utils.data import DataLoader\n",
-    "from torchvision import datasets, transforms\n",
-    "import numpy as np\n",
-    "import itertools\n",
-    "import matplotlib.pyplot as plt"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 1.2 Define network and SNN parameters\n",
-    "We will use a 784-1000-10 FCN architecture for a sequence of 25 time steps.\n",
-    "\n",
-    "* `alpha` is the decay rate of the synaptic current of a neuron\n",
-    "* `beta` is the decay rate of the membrane potential of a neuron"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# Network Architecture\n",
-    "num_inputs = 28*28\n",
-    "num_hidden = 1000\n",
-    "num_outputs = 10\n",
-    "\n",
-    "# Training Parameters\n",
-    "batch_size=128\n",
-    "data_path='/data/mnist'\n",
-    "\n",
-    "# Temporal Dynamics\n",
-    "num_steps = 25\n",
-    "time_step = 1e-3\n",
-    "tau_mem = 3e-3\n",
-    "tau_syn = 2.2e-3\n",
-    "alpha = float(np.exp(-time_step/tau_syn))\n",
-    "beta = float(np.exp(-time_step/tau_mem))\n",
-    "\n",
-    "dtype = torch.float\n",
-    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 1.3 Download MNIST Dataset\n",
-    "To see how to construct a validation set, refer to Tutorial 1."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# Define a transform\n",
-    "transform = transforms.Compose([\n",
-    "            transforms.Resize((28, 28)),\n",
-    "            transforms.Grayscale(),\n",
-    "            transforms.ToTensor(),\n",
-    "            transforms.Normalize((0,), (1,))])\n",
-    "\n",
-    "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n",
-    "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 1.4 Create DataLoaders"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
-    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 2. Define Network\n",
-    "snnTorch treats neurons as activations with recurrent connections. This allows for smooth integration with PyTorch.\n",
-    "There are a few useful neuron models and surrogate gradient functions which approximate the gradient of spikes.\n",
-    "\n",
-    "Our network will use one type of neuron model and one surrogate gradient:\n",
-    "1. `snntorch.Stein` is a basic leaky integrate and fire (LIF) neuron. Specifically, it assumes instantaneous rise times for synaptic current and membrane potential.\n",
-    "2. `snntorch.FastSigmoidSurrogate` defines separate forward and backward functions. The forward function is a Heaviside step function for spike generation. The backward function is the derivative of a fast sigmoid function, to ensure continuous differentiability.\n",
-    "The `FastSigmoidSurrogate` function has been adapted from:\n",
-    "\n",
-    ">Neftci, E. O., Mostafa, H., and Zenke, F. (2019) Surrogate Gradient Learning in Spiking Neural Networks. https://arxiv.org/abs/1901/09948"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# from snntorch import surrogate\n",
-    "#\n",
-    "# spike_grad = surrogate.FastSigmoid.apply\n",
-    "# snn.slope = 50"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "The surrogate is passed to `spike_grad` and overrides the default gradient of the Heaviside step function.\n",
-    "If we did not override the default gradient, (zero everywhere, except for $x=1$ where it is technically infinite but clipped to 1 here), then learning would not take place for as long as the neuron was not emitting post-synaptic spikes.\n",
-    "\n",
-    "`snn.slope` defines the slope of the backward surrogate.\n",
-    "\n",
-    "TO-DO: Include visualisation."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "Now we can define our spiking neural network (SNN).\n",
-    "Creating an instance of the `Stein` neuron requires two arguments and two optional arguments:\n",
-    "1. $I_{syn}$ decay rate, $\\alpha$,\n",
-    "2. $V_{mem}$ decay rate, $\\beta$,\n",
-    "3. the surrogate spiking function, `spike_grad` (*default*: the gradient of the Heaviside function), and\n",
-    "4. the threshold for spiking, (*default*: 1.0).\n",
-    "\n",
-    "snnTorch treats the LIF neuron as a recurrent activation. Therefore, it requires initialization of its internal states.\n",
-    "For each layer, we initialize the synaptic current `syn1` and `syn2`, the membrane potential `mem1` and `mem2`, and the post-synaptic spikes `spk1` and `spk2` to zero.\n",
-    "A class method `init_stein` will take care of this.\n",
-    "\n",
-    "For rate coding, the final layer of spikes and membrane potential are used to determine accuracy and loss, respectively.\n",
-    "So their historical values are recorded in `spk2_rec` and `mem2_rec`.\n",
-    "\n",
-    "Keep in mind, the dataset we are using is just static MNIST. I.e., it is *not* time-varying.\n",
-    "Therefore, we pass the same MNIST sample to the input at each time step.\n",
-    "This is handled in the line `cur1 = self.fc1(x)`, where `x` is the same input over the whole for-loop."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# Define Network\n",
-    "class Net(nn.Module):\n",
-    "    def __init__(self):\n",
-    "        super().__init__()\n",
-    "\n",
-    "    # initialize layers\n",
-    "        self.fc1 = nn.Linear(num_inputs, num_hidden)\n",
-    "        self.lif1 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
-    "        self.fc2 = nn.Linear(num_hidden, num_outputs)\n",
-    "        self.lif2 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
-    "\n",
-    "    def forward(self, x):\n",
-    "        spk1, syn1, mem1 = self.lif1.init_stein(batch_size, num_hidden)\n",
-    "        spk2, syn2, mem2 = self.lif2.init_stein(batch_size, num_outputs)\n",
-    "\n",
-    "        spk2_rec = []\n",
-    "        mem2_rec = []\n",
-    "\n",
-    "        for step in range(num_steps):\n",
-    "            cur1 = self.fc1(x)\n",
-    "            spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)\n",
-    "            cur2 = self.fc2(spk1)\n",
-    "            spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)\n",
-    "\n",
-    "            spk2_rec.append(spk2)\n",
-    "            mem2_rec.append(mem2)\n",
-    "\n",
-    "        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n",
-    "\n",
-    "net = Net().to(device)"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 3. Training\n",
-    "Time for training! Let's first define a couple of functions to print out test/train accuracy."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "outputs": [],
-   "source": [
-    "def print_batch_accuracy(data, targets, train=False):\n",
-    "    output, _ = net(data.view(batch_size, -1))\n",
-    "    _, idx = output.sum(dim=0).max(1)\n",
-    "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
-    "\n",
-    "    if train:\n",
-    "        print(f\"Train Set Accuracy: {acc}\")\n",
-    "    else:\n",
-    "        print(f\"Test Set Accuracy: {acc}\")\n",
-    "\n",
-    "def train_printer():\n",
-    "    print(f\"Epoch {epoch}, Minibatch {minibatch_counter}\")\n",
-    "    print(f\"Train Set Loss: {loss_hist[counter]}\")\n",
-    "    print(f\"Test Set Loss: {test_loss_hist[counter]}\")\n",
-    "    print_batch_accuracy(data_it, targets_it, train=True)\n",
-    "    print_batch_accuracy(testdata_it, testtargets_it, train=False)\n",
-    "    print(\"\\n\")"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 3.1 Optimizer & Loss\n",
-    "* *Output Activation*: We'll apply the softmax function to the membrane potentials of the output layer, rather than the spikes.\n",
-    "* *Loss*: This will then be used to calculate the negative log-likelihood loss.\n",
-    "By encouraging the membrane of the correct neuron class to reach the threshold, we expect that neuron will fire more frequently.\n",
-    "The loss could be applied to the spike count as well, but the membrane is  continuous whereas spike count is discrete.\n",
-    "* *Optimizer*: The Adam optimizer is used for weight updates.\n",
-    "* *Accuracy*: Accuracy is measured by counting the spikes of the output neurons. The neuron that fires the most frequently will be our predicted class."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "optimizer = torch.optim.Adam(net.parameters(), lr=2e-4, betas=(0.9, 0.999))\n",
-    "log_softmax_fn = nn.LogSoftmax(dim=-1)\n",
-    "loss_fn = nn.NLLLoss()"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 3.2 Training Loop\n",
-    "Now just sit back, relax, and wait for convergence."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "loss_hist = []\n",
-    "test_loss_hist = []\n",
-    "counter = 0\n",
-    "\n",
-    "# Outer training loop\n",
-    "for epoch in range(3):\n",
-    "    minibatch_counter = 0\n",
-    "    train_batch = iter(train_loader)\n",
-    "\n",
-    "    # Minibatch training loop\n",
-    "    for data_it, targets_it in train_batch:\n",
-    "        data_it = data_it.to(device)\n",
-    "        targets_it = targets_it.to(device)\n",
-    "\n",
-    "        output, mem_rec = net(data_it.view(batch_size, -1))\n",
-    "        log_p_y = log_softmax_fn(mem_rec)\n",
-    "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
-    "\n",
-    "        # Sum loss over time steps: BPTT\n",
-    "        for step in range(num_steps):\n",
-    "          loss_val += loss_fn(log_p_y[step], targets_it)\n",
-    "\n",
-    "        # Gradient calculation\n",
-    "        optimizer.zero_grad()\n",
-    "        loss_val.backward(retain_graph=True)\n",
-    "\n",
-    "        # Weight Update\n",
-    "        nn.utils.clip_grad_norm_(net.parameters(), 1) # gradient clipping\n",
-    "        optimizer.step()\n",
-    "\n",
-    "        # Store loss history for future plotting\n",
-    "        loss_hist.append(loss_val.item())\n",
-    "\n",
-    "        # Test set\n",
-    "        test_data = itertools.cycle(test_loader)\n",
-    "        testdata_it, testtargets_it = next(test_data)\n",
-    "        testdata_it = testdata_it.to(device)\n",
-    "        testtargets_it = testtargets_it.to(device)\n",
-    "\n",
-    "        # Test set forward pass\n",
-    "        test_output, test_mem_rec = net(testdata_it.view(batch_size, -1))\n",
-    "\n",
-    "        # Test set loss\n",
-    "        log_p_ytest = log_softmax_fn(test_mem_rec)\n",
-    "        log_p_ytest = log_p_ytest.sum(dim=0)\n",
-    "        loss_val_test = loss_fn(log_p_ytest, testtargets_it)\n",
-    "        test_loss_hist.append(loss_val_test.item())\n",
-    "\n",
-    "        # Print test/train loss/accuracy\n",
-    "        if counter % 50 == 0:\n",
-    "            train_printer()\n",
-    "        minibatch_counter += 1\n",
-    "        counter += 1\n",
-    "\n",
-    "loss_hist_true_grad = loss_hist\n",
-    "test_loss_hist_true_grad = test_loss_hist"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 4. Results\n",
-    "### 4.1 Plot Training/Test Loss"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# Plot Loss\n",
-    "fig = plt.figure(facecolor=\"w\", figsize=(10, 5))\n",
-    "plt.plot(loss_hist)\n",
-    "plt.plot(test_loss_hist)\n",
-    "plt.legend([\"Test Loss\", \"Train Loss\"])\n",
-    "plt.xlabel(\"Epoch\")\n",
-    "plt.ylabel(\"Loss\")\n",
-    "plt.show()"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "rtrNT4NPRp7r"
+      },
+      "source": [
+        "<img src='https://github.com/jeshraghian/snntorch/blob/master/docs/_static/img/snntorch_alpha_w.png?raw=true' width=\"400\">\n",
+        "\n",
+        "# snnTorch - Gradient-based Learning in Spiking Neural Networks\n",
+        "## Tutorial 2\n",
+        "### By Jason K. Eshraghian\n",
+        "\n",
+        "<a href=\"https://colab.research.google.com/github/googlecolab/colabtools/blob/master/notebooks/colab-github-demo.ipynb\">\n",
+        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+        "</a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "TKy-qDQdRp73"
+      },
+      "source": [
+        "# Introduction\n",
+        "In this tutorial, you will learn how to use snnTorch to:\n",
+        "* create a 2-layer fully-connected spiking network;\n",
+        "* implement the backpropagation through time (BPTT) algorithm;\n",
+        "* to classify both the static and spiking MNIST datasets.\n",
+        "\n",
+        "If running in Google Colab:\n",
+        "* You may connect to GPU by checking `Runtime` > `Change runtime type` > `Hardware accelerator: GPU`\n",
+        "* Next, install the latest PyPi distribution of snnTorch by clicking into the following cell and pressing `Shift+Enter`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "BgBRVUtpRp74"
+      },
+      "source": [
+        "!pip install snntorch"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Zm-D2lthRp75"
+      },
+      "source": [
+        "## 1. Setting up the Static MNIST Dataset\n",
+        "### 1.1. Import packages and setup environment"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "sEygpdc8Rp76"
+      },
+      "source": [
+        "import snntorch as snn\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torchvision import datasets, transforms\n",
+        "import numpy as np\n",
+        "import itertools\n",
+        "import matplotlib.pyplot as plt"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "RcX7J9vVRp76"
+      },
+      "source": [
+        "### 1.2 Define network and SNN parameters\n",
+        "We will use a 784-1000-10 FCN architecture for a sequence of 25 time steps.\n",
+        "\n",
+        "* `alpha` is the decay rate of the synaptic current of a neuron\n",
+        "* `beta` is the decay rate of the membrane potential of a neuron"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "bKEj2hucRp77"
+      },
+      "source": [
+        "# Network Architecture\n",
+        "num_inputs = 28*28\n",
+        "num_hidden = 1000\n",
+        "num_outputs = 10\n",
+        "\n",
+        "# Training Parameters\n",
+        "batch_size=128\n",
+        "data_path='/data/mnist'\n",
+        "\n",
+        "# Temporal Dynamics\n",
+        "num_steps = 25\n",
+        "time_step = 1e-3\n",
+        "tau_mem = 3e-3\n",
+        "tau_syn = 2.2e-3\n",
+        "alpha = float(np.exp(-time_step/tau_syn))\n",
+        "beta = float(np.exp(-time_step/tau_mem))\n",
+        "\n",
+        "dtype = torch.float\n",
+        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "gM_hcwDIRp78"
+      },
+      "source": [
+        "### 1.3 Download MNIST Dataset\n",
+        "To see how to construct a validation set, refer to Tutorial 1."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "0xwYb15xRp79"
+      },
+      "source": [
+        "# Define a transform\n",
+        "transform = transforms.Compose([\n",
+        "            transforms.Resize((28, 28)),\n",
+        "            transforms.Grayscale(),\n",
+        "            transforms.ToTensor(),\n",
+        "            transforms.Normalize((0,), (1,))])\n",
+        "\n",
+        "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n",
+        "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "DIg3vLmURp7-"
+      },
+      "source": [
+        "### 1.4 Create DataLoaders"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "K4DF-odMRp7_"
+      },
+      "source": [
+        "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
+        "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "pD4Dw-RoRp7_"
+      },
+      "source": [
+        "## 2. Define Network\n",
+        "snnTorch treats neurons as activations with recurrent connections. This allows for smooth integration with PyTorch.\n",
+        "There are a few useful neuron models and surrogate gradient functions which approximate the gradient of spikes.\n",
+        "\n",
+        "Our network will use one type of neuron model and one surrogate gradient:\n",
+        "1. `snntorch.Stein` is a basic leaky integrate and fire (LIF) neuron. Specifically, it assumes instantaneous rise times for synaptic current and membrane potential.\n",
+        "2. `snntorch.FastSigmoidSurrogate` defines separate forward and backward functions. The forward function is a Heaviside step function for spike generation. The backward function is the derivative of a fast sigmoid function, to ensure continuous differentiability.\n",
+        "The `FastSigmoidSurrogate` function has been adapted from:\n",
+        "\n",
+        ">Neftci, E. O., Mostafa, H., and Zenke, F. (2019) Surrogate Gradient Learning in Spiking Neural Networks. https://arxiv.org/abs/1901/09948"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "Sf9RdE9jRp8A"
+      },
+      "source": [
+        "# from snntorch import surrogate\n",
+        "#\n",
+        "# spike_grad = surrogate.FastSigmoid.apply\n",
+        "# snn.slope = 50"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Og9e57W0Rp8B"
+      },
+      "source": [
+        "The surrogate is passed to `spike_grad` and overrides the default gradient of the Heaviside step function.\n",
+        "If we did not override the default gradient, (zero everywhere, except for $x=1$ where it is technically infinite but clipped to 1 here), then learning would not take place for as long as the neuron was not emitting post-synaptic spikes.\n",
+        "\n",
+        "`snn.slope` defines the slope of the backward surrogate.\n",
+        "\n",
+        "TO-DO: Include visualisation."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Yo-x48mARp8C"
+      },
+      "source": [
+        "Now we can define our spiking neural network (SNN).\n",
+        "Creating an instance of the `Stein` neuron requires two arguments and two optional arguments:\n",
+        "1. $I_{syn}$ decay rate, $\\alpha$,\n",
+        "2. $V_{mem}$ decay rate, $\\beta$,\n",
+        "3. the surrogate spiking function, `spike_grad` (*default*: the gradient of the Heaviside function), and\n",
+        "4. the threshold for spiking, (*default*: 1.0).\n",
+        "\n",
+        "snnTorch treats the LIF neuron as a recurrent activation. Therefore, it requires initialization of its internal states.\n",
+        "For each layer, we initialize the synaptic current `syn1` and `syn2`, the membrane potential `mem1` and `mem2`, and the post-synaptic spikes `spk1` and `spk2` to zero.\n",
+        "A class method `init_stein` will take care of this.\n",
+        "\n",
+        "For rate coding, the final layer of spikes and membrane potential are used to determine accuracy and loss, respectively.\n",
+        "So their historical values are recorded in `spk2_rec` and `mem2_rec`.\n",
+        "\n",
+        "Keep in mind, the dataset we are using is just static MNIST. I.e., it is *not* time-varying.\n",
+        "Therefore, we pass the same MNIST sample to the input at each time step.\n",
+        "This is handled in the line `cur1 = self.fc1(x)`, where `x` is the same input over the whole for-loop."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "P6RHCnXMRp8D"
+      },
+      "source": [
+        "# Define Network\n",
+        "class Net(nn.Module):\n",
+        "    def __init__(self):\n",
+        "        super().__init__()\n",
+        "\n",
+        "    # initialize layers\n",
+        "        self.fc1 = nn.Linear(num_inputs, num_hidden)\n",
+        "        self.lif1 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "        self.fc2 = nn.Linear(num_hidden, num_outputs)\n",
+        "        self.lif2 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        spk1, syn1, mem1 = self.lif1.init_stein(batch_size, num_hidden)\n",
+        "        spk2, syn2, mem2 = self.lif2.init_stein(batch_size, num_outputs)\n",
+        "\n",
+        "        spk2_rec = []\n",
+        "        mem2_rec = []\n",
+        "\n",
+        "        for step in range(num_steps):\n",
+        "            cur1 = self.fc1(x)\n",
+        "            spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)\n",
+        "            cur2 = self.fc2(spk1)\n",
+        "            spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)\n",
+        "\n",
+        "            spk2_rec.append(spk2)\n",
+        "            mem2_rec.append(mem2)\n",
+        "\n",
+        "        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n",
+        "\n",
+        "net = Net().to(device)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "386KNHG7Rp8E"
+      },
+      "source": [
+        "## 3. Training\n",
+        "Time for training! Let's first define a couple of functions to print out test/train accuracy."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "cOKKbUnDRp8F"
+      },
+      "source": [
+        "def print_batch_accuracy(data, targets, train=False):\n",
+        "    output, _ = net(data.view(batch_size, -1))\n",
+        "    _, idx = output.sum(dim=0).max(1)\n",
+        "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
+        "\n",
+        "    if train:\n",
+        "        print(f\"Train Set Accuracy: {acc}\")\n",
+        "    else:\n",
+        "        print(f\"Test Set Accuracy: {acc}\")\n",
+        "\n",
+        "def train_printer():\n",
+        "    print(f\"Epoch {epoch}, Minibatch {minibatch_counter}\")\n",
+        "    print(f\"Train Set Loss: {loss_hist[counter]}\")\n",
+        "    print(f\"Test Set Loss: {test_loss_hist[counter]}\")\n",
+        "    print_batch_accuracy(data_it, targets_it, train=True)\n",
+        "    print_batch_accuracy(testdata_it, testtargets_it, train=False)\n",
+        "    print(\"\\n\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "mZqfCe0KRp8J"
+      },
+      "source": [
+        "### 3.1 Optimizer & Loss\n",
+        "* *Output Activation*: We'll apply the softmax function to the membrane potentials of the output layer, rather than the spikes.\n",
+        "* *Loss*: This will then be used to calculate the negative log-likelihood loss.\n",
+        "By encouraging the membrane of the correct neuron class to reach the threshold, we expect that neuron will fire more frequently.\n",
+        "The loss could be applied to the spike count as well, but the membrane is  continuous whereas spike count is discrete.\n",
+        "* *Optimizer*: The Adam optimizer is used for weight updates.\n",
+        "* *Accuracy*: Accuracy is measured by counting the spikes of the output neurons. The neuron that fires the most frequently will be our predicted class."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "UMZ-4uGlRp8K"
+      },
+      "source": [
+        "optimizer = torch.optim.Adam(net.parameters(), lr=2e-4, betas=(0.9, 0.999))\n",
+        "log_softmax_fn = nn.LogSoftmax(dim=-1)\n",
+        "loss_fn = nn.NLLLoss()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "UerqLkd1Rp8K"
+      },
+      "source": [
+        "### 3.2 Training Loop\n",
+        "Now just sit back, relax, and wait for convergence."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "76CEHI2xRp8L"
+      },
+      "source": [
+        "loss_hist = []\n",
+        "test_loss_hist = []\n",
+        "counter = 0\n",
+        "\n",
+        "# Outer training loop\n",
+        "for epoch in range(3):\n",
+        "    minibatch_counter = 0\n",
+        "    train_batch = iter(train_loader)\n",
+        "\n",
+        "    # Minibatch training loop\n",
+        "    for data_it, targets_it in train_batch:\n",
+        "        data_it = data_it.to(device)\n",
+        "        targets_it = targets_it.to(device)\n",
+        "\n",
+        "        output, mem_rec = net(data_it.view(batch_size, -1))\n",
+        "        log_p_y = log_softmax_fn(mem_rec)\n",
+        "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
+        "\n",
+        "        # Sum loss over time steps: BPTT\n",
+        "        for step in range(num_steps):\n",
+        "          loss_val += loss_fn(log_p_y[step], targets_it)\n",
+        "\n",
+        "        # Gradient calculation\n",
+        "        optimizer.zero_grad()\n",
+        "        loss_val.backward(retain_graph=True)\n",
+        "\n",
+        "        # Weight Update\n",
+        "        nn.utils.clip_grad_norm_(net.parameters(), 1) # gradient clipping\n",
+        "        optimizer.step()\n",
+        "\n",
+        "        # Store loss history for future plotting\n",
+        "        loss_hist.append(loss_val.item())\n",
+        "\n",
+        "        # Test set\n",
+        "        test_data = itertools.cycle(test_loader)\n",
+        "        testdata_it, testtargets_it = next(test_data)\n",
+        "        testdata_it = testdata_it.to(device)\n",
+        "        testtargets_it = testtargets_it.to(device)\n",
+        "\n",
+        "        # Test set forward pass\n",
+        "        test_output, test_mem_rec = net(testdata_it.view(batch_size, -1))\n",
+        "\n",
+        "        # Test set loss\n",
+        "        log_p_ytest = log_softmax_fn(test_mem_rec)\n",
+        "        log_p_ytest = log_p_ytest.sum(dim=0)\n",
+        "        loss_val_test = loss_fn(log_p_ytest, testtargets_it)\n",
+        "        test_loss_hist.append(loss_val_test.item())\n",
+        "\n",
+        "        # Print test/train loss/accuracy\n",
+        "        if counter % 50 == 0:\n",
+        "            train_printer()\n",
+        "        minibatch_counter += 1\n",
+        "        counter += 1\n",
+        "\n",
+        "loss_hist_true_grad = loss_hist\n",
+        "test_loss_hist_true_grad = test_loss_hist"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "fdOuwCWHRp8L"
+      },
+      "source": [
+        "## 4. Results\n",
+        "### 4.1 Plot Training/Test Loss"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "WJGSBq6zRp8M"
+      },
+      "source": [
+        "# Plot Loss\n",
+        "fig = plt.figure(facecolor=\"w\", figsize=(10, 5))\n",
+        "plt.plot(loss_hist)\n",
+        "plt.plot(test_loss_hist)\n",
+        "plt.legend([\"Test Loss\", \"Train Loss\"])\n",
+        "plt.xlabel(\"Epoch\")\n",
+        "plt.ylabel(\"Loss\")\n",
+        "plt.show()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "h7xb37iHRp8N"
+      },
+      "source": [
+        "### 4.2 Test Set Accuracy\n",
+        "This function just iterates over all minibatches to obtain a measure of accuracy over the full 10,000 samples in the test set."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "R1ReGuNURp8N"
+      },
+      "source": [
+        "total = 0\n",
+        "correct = 0\n",
+        "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "  net.eval()\n",
+        "  for data in test_loader:\n",
+        "    images, labels = data\n",
+        "    images = images.to(device)\n",
+        "    labels = labels.to(device)\n",
+        "\n",
+        "    # If current batch matches batch_size, just do the usual thing\n",
+        "    if images.size()[0] == batch_size:\n",
+        "      outputs, _ = net(images.view(batch_size, -1))\n",
+        "\n",
+        "    # If current batch does not match batch_size (e.g., is the final minibatch),\n",
+        "    # modify batch_size in a temp variable and restore it at the end of the else block\n",
+        "    else:\n",
+        "      temp_bs = batch_size\n",
+        "      batch_size = images.size()[0]\n",
+        "      outputs, _ = net(images.view(images.size()[0], -1))\n",
+        "      batch_size = temp_bs\n",
+        "\n",
+        "    _, predicted = outputs.sum(dim=0).max(1)\n",
+        "    total += labels.size(0)\n",
+        "    correct += (predicted == labels).sum().item()\n",
+        "\n",
+        "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
+        "print(f\"Test Set Accuracy: {100 * correct / total}%\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "zNJjSKATRp8N"
+      },
+      "source": [
+        "Voila! That's it for static MNIST."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "2qUxjHwBRp8f"
+      },
+      "source": [
+        "## 5. Spiking MNIST\n",
+        "Part of the appeal of SNNs is their ability to handle time-varying spiking data. So let's use rate-coding to convert MNIST into spiking MNIST using the `spikegen` module in the previous tutorial, and train our network with that instead."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "8K7_C-2rRp8g"
+      },
+      "source": [
+        "from snntorch import spikegen\n",
+        "\n",
+        "# MNIST to spiking-MNIST\n",
+        "spike_data, spike_targets = spikegen.rate(data_it, targets_it, num_outputs=num_outputs, num_steps=num_steps, gain=1,\n",
+        "                                          offset=0, convert_targets=False, temporal_targets=False)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "BvaNfns_Rp8h"
+      },
+      "source": [
+        "### 5.1 Visualiser\n",
+        "Just so you're damn sure it's a spiking input."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "ll2D4jtdRp8i"
+      },
+      "source": [
+        "!pip install celluloid # matplotlib animations made easy"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "cJQ0XUAZRp8i"
+      },
+      "source": [
+        "from celluloid import Camera\n",
+        "from IPython.display import HTML\n",
+        "\n",
+        "# Animator\n",
+        "spike_data_sample = spike_data[:, 0, 0].cpu()\n",
+        "\n",
+        "fig, ax = plt.subplots()\n",
+        "camera = Camera(fig)\n",
+        "plt.axis('off')\n",
+        "\n",
+        "for step in range(num_steps):\n",
+        "    im = ax.imshow(spike_data_sample[step, :, :], cmap='plasma')\n",
+        "    camera.snap()\n",
+        "\n",
+        "# interval=40 specifies 40ms delay between frames\n",
+        "a = camera.animate(interval=40)\n",
+        "HTML(a.to_html5_video())"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "T-dMHLBbRp8k"
+      },
+      "source": [
+        "print(spike_targets[0])"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "I63CROTaRp8k"
+      },
+      "source": [
+        "## 6. Define Network\n",
+        "The network is the same as before. The one difference is that the for-loop iterates through the first dimension of the input:\n",
+        "`cur1 = self.fc1(x[step])`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "a_elvRYIRp8k"
+      },
+      "source": [
+        "spike_grad = snn.FastSigmoidSurrogate.apply\n",
+        "snn.slope = 50 # The lower the slope, the smoother the gradient\n",
+        "\n",
+        "# Define Network\n",
+        "class Net(nn.Module):\n",
+        "    def __init__(self):\n",
+        "        super().__init__()\n",
+        "\n",
+        "        # Initialize layers\n",
+        "        self.fc1 = nn.Linear(num_inputs, num_hidden)\n",
+        "        self.lif1 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "        self.fc2 = nn.Linear(num_hidden, num_outputs)\n",
+        "        self.lif2 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        # Initialize hidden states + output spike at t=0\n",
+        "        spk1, syn1, mem1 = self.lif1.init_stein(batch_size, num_hidden)\n",
+        "        spk2, syn2, mem2 = self.lif2.init_stein(batch_size, num_outputs)\n",
+        "\n",
+        "        spk2_rec = []\n",
+        "        mem2_rec = []\n",
+        "\n",
+        "        for step in range(num_steps):\n",
+        "            cur1 = self.fc1(x[step])\n",
+        "            spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)\n",
+        "            cur2 = self.fc2(spk1)\n",
+        "            spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)\n",
+        "\n",
+        "            spk2_rec.append(spk2)\n",
+        "            mem2_rec.append(mem2)\n",
+        "\n",
+        "        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n",
+        "\n",
+        "net = Net().to(device)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "E8utR7I9Rp8l"
+      },
+      "source": [
+        "## 7. Training\n",
+        "We make a slight modification to our print-out functions to handle the new first dimension of the input:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "FlkQt3skRp8l"
+      },
+      "source": [
+        "def print_batch_accuracy(data, targets, train=False):\n",
+        "    output, _ = net(data.view(num_steps, batch_size, -1))\n",
+        "    _, idx = output.sum(dim=0).max(1)\n",
+        "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
+        "\n",
+        "    if train:\n",
+        "        print(f\"Train Set Accuracy: {acc}\")\n",
+        "    else:\n",
+        "        print(f\"Test Set Accuracy: {acc}\")\n",
+        "\n",
+        "def train_printer():\n",
+        "    print(f\"Epoch {epoch}, Minibatch {minibatch_counter}\")\n",
+        "    print(f\"Train Set Loss: {loss_hist[counter]}\")\n",
+        "    print(f\"Test Set Loss: {test_loss_hist[counter]}\")\n",
+        "    print_batch_accuracy(spike_data, spike_targets, train=True)\n",
+        "    print_batch_accuracy(test_spike_data, test_spike_targets, train=False)\n",
+        "    print(\"\\n\")\n"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "vbve35sDRp8l"
+      },
+      "source": [
+        "### 7.1 Optimizer & Loss\n",
+        "We'll keep our optimizer and loss the exact same as the static MNIST case."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "10CH8RKHRp8l"
+      },
+      "source": [
+        "optimizer = torch.optim.Adam(net.parameters(), lr=2e-4, betas=(0.9, 0.999))\n",
+        "log_softmax_fn = nn.LogSoftmax(dim=-1)\n",
+        "loss_fn = nn.NLLLoss()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Ef0TM6yIRp8m"
+      },
+      "source": [
+        "### 7.2 Training Loop\n",
+        "The training loop is identical to the static MNIST case, but we pass each minibatch through `spikegen.rate` before running it through the feedforward network."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "ISLzPpwvRp8m"
+      },
+      "source": [
+        "loss_hist = []\n",
+        "test_loss_hist = []\n",
+        "counter = 0\n",
+        "\n",
+        "# Outer training loop\n",
+        "for epoch in range(3):\n",
+        "    minibatch_counter = 0\n",
+        "    data = iter(train_loader)\n",
+        "\n",
+        "    # Minibatch training loop\n",
+        "    for data_it, targets_it in data:\n",
+        "        data_it = data_it.to(device)\n",
+        "        targets_it = targets_it.to(device)\n",
+        "\n",
+        "        # Spike generator\n",
+        "        spike_data, spike_targets = spikegen.rate(data_it, targets_it, num_outputs=num_outputs, num_steps=num_steps,\n",
+        "                                                  gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "        # Forward pass\n",
+        "        output, mem_rec = net(spike_data.view(num_steps, batch_size, -1))\n",
+        "        log_p_y = log_softmax_fn(mem_rec)\n",
+        "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
+        "\n",
+        "        # Sum loss over time steps to perform BPTT\n",
+        "        for step in range(num_steps):\n",
+        "          loss_val += loss_fn(log_p_y[step], targets_it)\n",
+        "\n",
+        "        # Gradient Calculation\n",
+        "        optimizer.zero_grad()\n",
+        "        loss_val.backward(retain_graph=True)\n",
+        "        nn.utils.clip_grad_norm_(net.parameters(), 1)\n",
+        "\n",
+        "        # Weight Update\n",
+        "        optimizer.step()\n",
+        "\n",
+        "        # Store Loss history\n",
+        "        loss_hist.append(loss_val.item())\n",
+        "\n",
+        "        # Test set\n",
+        "        test_data = itertools.cycle(test_loader)\n",
+        "        testdata_it, testtargets_it = next(test_data)\n",
+        "        testdata_it = testdata_it.to(device)\n",
+        "        testtargets_it = testtargets_it.to(device)\n",
+        "\n",
+        "        # Test set spike conversion\n",
+        "        test_spike_data, test_spike_targets = spikegen.rate(testdata_it, testtargets_it, num_outputs=num_outputs,\n",
+        "                                                            num_steps=num_steps, gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "        # Test set forward pass\n",
+        "        test_output, test_mem_rec = net(test_spike_data.view(num_steps, batch_size, -1))\n",
+        "\n",
+        "        # Test set loss\n",
+        "        log_p_ytest = log_softmax_fn(test_mem_rec)\n",
+        "        log_p_ytest = log_p_ytest.sum(dim=0)\n",
+        "        loss_val_test = loss_fn(log_p_ytest, test_spike_targets)\n",
+        "        test_loss_hist.append(loss_val_test.item())\n",
+        "\n",
+        "        # Print test/train loss/accuracy\n",
+        "        if counter % 50 == 0:\n",
+        "            train_printer()\n",
+        "        minibatch_counter += 1\n",
+        "        counter += 1\n",
+        "\n",
+        "loss_hist_true_grad = loss_hist\n",
+        "test_loss_hist_true_grad = test_loss_hist"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "EuRL61AnRp8m"
+      },
+      "source": [
+        "## 8. Spiking MNIST Results\n",
+        "### 8.1 Plot Training/Test Loss"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "fl5griHBRp8n"
+      },
+      "source": [
+        "# Plot Loss\n",
+        "fig = plt.figure(facecolor=\"w\", figsize=(10, 5))\n",
+        "plt.plot(loss_hist)\n",
+        "plt.plot(test_loss_hist)\n",
+        "plt.legend([\"Test Loss\", \"Train Loss\"])\n",
+        "plt.xlabel(\"Epoch\")\n",
+        "plt.ylabel(\"Loss\")\n",
+        "plt.show()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "IZMH5UOlRp8n"
+      },
+      "source": [
+        "### 8.2 Test Set Accuracy"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "eO2Krz1mRp8n"
+      },
+      "source": [
+        "total = 0\n",
+        "correct = 0\n",
+        "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "  net.eval()\n",
+        "  for data in test_loader:\n",
+        "    images, labels = data\n",
+        "    images = images.to(device)\n",
+        "    labels = labels.to(device)\n",
+        "\n",
+        "    # If current batch matches batch_size, just do the usual thing\n",
+        "    if images.size()[0] == batch_size:\n",
+        "      spike_test, spike_targets = spikegen.rate(images, labels, num_outputs=num_outputs, num_steps=num_steps,\n",
+        "                                                            gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "      outputs, _ = net(spike_test.view(num_steps, batch_size, -1))\n",
+        "\n",
+        "    # If current batch does not match batch_size (e.g., is the final minibatch),\n",
+        "    # modify batch_size in a temp variable and restore it at the end of the else block\n",
+        "    else:\n",
+        "      temp_bs = batch_size\n",
+        "      batch_size = images.size()[0]\n",
+        "      spike_test, spike_targets = spikegen.rate(images, labels, num_outputs=num_outputs, num_steps=num_steps,\n",
+        "                                                gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "      outputs, _ = net(spike_test.view(num_steps, images.size()[0], -1))\n",
+        "      batch_size = temp_bs\n",
+        "\n",
+        "    _, predicted = outputs.sum(dim=0).max(1)\n",
+        "    total += spike_targets.size(0)\n",
+        "    correct += (predicted == spike_targets).sum().item()\n",
+        "\n",
+        "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
+        "print(f\"Test Set Accuracy: {100 * correct / total}%\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "ztM4ogAqRp8n"
+      },
+      "source": [
+        "That's all for now!\n",
+        "Next time, we'll introduce how to use spiking convolutional layers to improve accuracy."
+      ]
     }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 4.2 Test Set Accuracy\n",
-    "This function just iterates over all minibatches to obtain a measure of accuracy over the full 10,000 samples in the test set."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "total = 0\n",
-    "correct = 0\n",
-    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
-    "\n",
-    "with torch.no_grad():\n",
-    "  net.eval()\n",
-    "  for data in test_loader:\n",
-    "    images, labels = data\n",
-    "    images = images.to(device)\n",
-    "    labels = labels.to(device)\n",
-    "\n",
-    "    # If current batch matches batch_size, just do the usual thing\n",
-    "    if images.size()[0] == batch_size:\n",
-    "      outputs, _ = net(images.view(batch_size, -1))\n",
-    "\n",
-    "    # If current batch does not match batch_size (e.g., is the final minibatch),\n",
-    "    # modify batch_size in a temp variable and restore it at the end of the else block\n",
-    "    else:\n",
-    "      temp_bs = batch_size\n",
-    "      batch_size = images.size()[0]\n",
-    "      outputs, _ = net(images.view(images.size()[0], -1))\n",
-    "      batch_size = temp_bs\n",
-    "\n",
-    "    _, predicted = outputs.sum(dim=0).max(1)\n",
-    "    total += labels.size(0)\n",
-    "    correct += (predicted == labels).sum().item()\n",
-    "\n",
-    "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
-    "print(f\"Test Set Accuracy: {100 * correct / total}%\")"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "Voila! That's it for static MNIST."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 5. Spiking MNIST\n",
-    "Part of the appeal of SNNs is their ability to handle time-varying spiking data. So let's use rate-coding to convert MNIST into spiking MNIST using the `spikegen` module in the previous tutorial, and train our network with that instead."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "from snntorch import spikegen\n",
-    "\n",
-    "# MNIST to spiking-MNIST\n",
-    "spike_data, spike_targets = spikegen.rate(data_it, targets_it, num_outputs=num_outputs, num_steps=num_steps, gain=1,\n",
-    "                                          offset=0, convert_targets=False, temporal_targets=False)"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 5.1 Visualiser\n",
-    "Just so you're damn sure it's a spiking input."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "source": [
-    "!pip install celluloid # matplotlib animations made easy"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "execution_count": null,
-   "outputs": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "from celluloid import Camera\n",
-    "from IPython.display import HTML\n",
-    "\n",
-    "# Animator\n",
-    "spike_data_sample = spike_data[:, 0, 0].cpu()\n",
-    "\n",
-    "fig, ax = plt.subplots()\n",
-    "camera = Camera(fig)\n",
-    "plt.axis('off')\n",
-    "\n",
-    "for step in range(num_steps):\n",
-    "    im = ax.imshow(spike_data_sample[step, :, :], cmap='plasma')\n",
-    "    camera.snap()\n",
-    "\n",
-    "# interval=40 specifies 40ms delay between frames\n",
-    "a = camera.animate(interval=40)\n",
-    "HTML(a.to_html5_video())"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "print(spike_targets[0])"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 6. Define Network\n",
-    "The network is the same as before. The one difference is that the for-loop iterates through the first dimension of the input:\n",
-    "`cur1 = self.fc1(x[step])`"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "spike_grad = snn.FastSigmoidSurrogate.apply\n",
-    "snn.slope = 50 # The lower the slope, the smoother the gradient\n",
-    "\n",
-    "# Define Network\n",
-    "class Net(nn.Module):\n",
-    "    def __init__(self):\n",
-    "        super().__init__()\n",
-    "\n",
-    "        # Initialize layers\n",
-    "        self.fc1 = nn.Linear(num_inputs, num_hidden)\n",
-    "        self.lif1 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
-    "        self.fc2 = nn.Linear(num_hidden, num_outputs)\n",
-    "        self.lif2 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
-    "\n",
-    "    def forward(self, x):\n",
-    "        # Initialize hidden states + output spike at t=0\n",
-    "        spk1, syn1, mem1 = self.lif1.init_stein(batch_size, num_hidden)\n",
-    "        spk2, syn2, mem2 = self.lif2.init_stein(batch_size, num_outputs)\n",
-    "\n",
-    "        spk2_rec = []\n",
-    "        mem2_rec = []\n",
-    "\n",
-    "        for step in range(num_steps):\n",
-    "            cur1 = self.fc1(x[step])\n",
-    "            spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)\n",
-    "            cur2 = self.fc2(spk1)\n",
-    "            spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)\n",
-    "\n",
-    "            spk2_rec.append(spk2)\n",
-    "            mem2_rec.append(mem2)\n",
-    "\n",
-    "        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n",
-    "\n",
-    "net = Net().to(device)"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 7. Training\n",
-    "We make a slight modification to our print-out functions to handle the new first dimension of the input:"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "def print_batch_accuracy(data, targets, train=False):\n",
-    "    output, _ = net(data.view(num_steps, batch_size, -1))\n",
-    "    _, idx = output.sum(dim=0).max(1)\n",
-    "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
-    "\n",
-    "    if train:\n",
-    "        print(f\"Train Set Accuracy: {acc}\")\n",
-    "    else:\n",
-    "        print(f\"Test Set Accuracy: {acc}\")\n",
-    "\n",
-    "def train_printer():\n",
-    "    print(f\"Epoch {epoch}, Minibatch {minibatch_counter}\")\n",
-    "    print(f\"Train Set Loss: {loss_hist[counter]}\")\n",
-    "    print(f\"Test Set Loss: {test_loss_hist[counter]}\")\n",
-    "    print_batch_accuracy(spike_data, spike_targets, train=True)\n",
-    "    print_batch_accuracy(test_spike_data, test_spike_targets, train=False)\n",
-    "    print(\"\\n\")\n"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 7.1 Optimizer & Loss\n",
-    "We'll keep our optimizer and loss the exact same as the static MNIST case."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "optimizer = torch.optim.Adam(net.parameters(), lr=2e-4, betas=(0.9, 0.999))\n",
-    "log_softmax_fn = nn.LogSoftmax(dim=-1)\n",
-    "loss_fn = nn.NLLLoss()"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 7.2 Training Loop\n",
-    "The training loop is identical to the static MNIST case, but we pass each minibatch through `spikegen.rate` before running it through the feedforward network."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "loss_hist = []\n",
-    "test_loss_hist = []\n",
-    "counter = 0\n",
-    "\n",
-    "# Outer training loop\n",
-    "for epoch in range(3):\n",
-    "    minibatch_counter = 0\n",
-    "    data = iter(train_loader)\n",
-    "\n",
-    "    # Minibatch training loop\n",
-    "    for data_it, targets_it in data:\n",
-    "        data_it = data_it.to(device)\n",
-    "        targets_it = targets_it.to(device)\n",
-    "\n",
-    "        # Spike generator\n",
-    "        spike_data, spike_targets = spikegen.rate(data_it, targets_it, num_outputs=num_outputs, num_steps=num_steps,\n",
-    "                                                  gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
-    "\n",
-    "        # Forward pass\n",
-    "        output, mem_rec = net(spike_data.view(num_steps, batch_size, -1))\n",
-    "        log_p_y = log_softmax_fn(mem_rec)\n",
-    "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
-    "\n",
-    "        # Sum loss over time steps to perform BPTT\n",
-    "        for step in range(num_steps):\n",
-    "          loss_val += loss_fn(log_p_y[step], targets_it)\n",
-    "\n",
-    "        # Gradient Calculation\n",
-    "        optimizer.zero_grad()\n",
-    "        loss_val.backward(retain_graph=True)\n",
-    "        nn.utils.clip_grad_norm_(net.parameters(), 1)\n",
-    "\n",
-    "        # Weight Update\n",
-    "        optimizer.step()\n",
-    "\n",
-    "        # Store Loss history\n",
-    "        loss_hist.append(loss_val.item())\n",
-    "\n",
-    "        # Test set\n",
-    "        test_data = itertools.cycle(test_loader)\n",
-    "        testdata_it, testtargets_it = next(test_data)\n",
-    "        testdata_it = testdata_it.to(device)\n",
-    "        testtargets_it = testtargets_it.to(device)\n",
-    "\n",
-    "        # Test set spike conversion\n",
-    "        test_spike_data, test_spike_targets = spikegen.rate(testdata_it, testtargets_it, num_outputs=num_outputs,\n",
-    "                                                            num_steps=num_steps, gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
-    "\n",
-    "        # Test set forward pass\n",
-    "        test_output, test_mem_rec = net(test_spike_data.view(num_steps, batch_size, -1))\n",
-    "\n",
-    "        # Test set loss\n",
-    "        log_p_ytest = log_softmax_fn(test_mem_rec)\n",
-    "        log_p_ytest = log_p_ytest.sum(dim=0)\n",
-    "        loss_val_test = loss_fn(log_p_ytest, test_spike_targets)\n",
-    "        test_loss_hist.append(loss_val_test.item())\n",
-    "\n",
-    "        # Print test/train loss/accuracy\n",
-    "        if counter % 50 == 0:\n",
-    "            train_printer()\n",
-    "        minibatch_counter += 1\n",
-    "        counter += 1\n",
-    "\n",
-    "loss_hist_true_grad = loss_hist\n",
-    "test_loss_hist_true_grad = test_loss_hist"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## 8. Spiking MNIST Results\n",
-    "### 8.1 Plot Training/Test Loss"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "# Plot Loss\n",
-    "fig = plt.figure(facecolor=\"w\", figsize=(10, 5))\n",
-    "plt.plot(loss_hist)\n",
-    "plt.plot(test_loss_hist)\n",
-    "plt.legend([\"Test Loss\", \"Train Loss\"])\n",
-    "plt.xlabel(\"Epoch\")\n",
-    "plt.ylabel(\"Loss\")\n",
-    "plt.show()"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### 8.2 Test Set Accuracy"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "total = 0\n",
-    "correct = 0\n",
-    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
-    "\n",
-    "with torch.no_grad():\n",
-    "  net.eval()\n",
-    "  for data in test_loader:\n",
-    "    images, labels = data\n",
-    "    images = images.to(device)\n",
-    "    labels = labels.to(device)\n",
-    "\n",
-    "    # If current batch matches batch_size, just do the usual thing\n",
-    "    if images.size()[0] == batch_size:\n",
-    "      spike_test, spike_targets = spikegen.rate(images, labels, num_outputs=num_outputs, num_steps=num_steps,\n",
-    "                                                            gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
-    "\n",
-    "      outputs, _ = net(spike_test.view(num_steps, batch_size, -1))\n",
-    "\n",
-    "    # If current batch does not match batch_size (e.g., is the final minibatch),\n",
-    "    # modify batch_size in a temp variable and restore it at the end of the else block\n",
-    "    else:\n",
-    "      temp_bs = batch_size\n",
-    "      batch_size = images.size()[0]\n",
-    "      spike_test, spike_targets = spikegen.rate(images, labels, num_outputs=num_outputs, num_steps=num_steps,\n",
-    "                                                gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
-    "\n",
-    "      outputs, _ = net(spike_test.view(num_steps, images.size()[0], -1))\n",
-    "      batch_size = temp_bs\n",
-    "\n",
-    "    _, predicted = outputs.sum(dim=0).max(1)\n",
-    "    total += spike_targets.size(0)\n",
-    "    correct += (predicted == spike_targets).sum().item()\n",
-    "\n",
-    "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
-    "print(f\"Test Set Accuracy: {100 * correct / total}%\")"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "That's all for now!\n",
-    "Next time, we'll introduce how to use spiking convolutional layers to improve accuracy."
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "name": "python3",
-   "language": "python",
-   "display_name": "Python 3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 2
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.6"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 0
+  ]
 }
\ No newline at end of file
diff --git a/examples/tutorial_2_neurons.ipynb b/examples/tutorial_2_neurons.ipynb
new file mode 100644
index 00000000..6ce4ad96
--- /dev/null
+++ b/examples/tutorial_2_neurons.ipynb
@@ -0,0 +1,1055 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "kernelspec": {
+      "name": "python3",
+      "language": "python",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 2
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython2",
+      "version": "2.7.6"
+    },
+    "colab": {
+      "name": "tutorial_2_FCN.ipynb",
+      "provenance": []
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "rtrNT4NPRp7r"
+      },
+      "source": [
+        "<img src='https://github.com/jeshraghian/snntorch/blob/master/docs/_static/img/snntorch_alpha_w.png?raw=true' width=\"400\">\n",
+        "\n",
+        "# snnTorch - Gradient-based Learning in Spiking Neural Networks\n",
+        "## Tutorial 2\n",
+        "### By Jason K. Eshraghian\n",
+        "\n",
+        "<a href=\"https://colab.research.google.com/github/googlecolab/colabtools/blob/master/notebooks/colab-github-demo.ipynb\">\n",
+        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+        "</a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "TKy-qDQdRp73"
+      },
+      "source": [
+        "# Introduction\n",
+        "In this tutorial, you will learn how to use snnTorch to:\n",
+        "* create a 2-layer fully-connected spiking network;\n",
+        "* implement the backpropagation through time (BPTT) algorithm;\n",
+        "* to classify both the static and spiking MNIST datasets.\n",
+        "\n",
+        "If running in Google Colab:\n",
+        "* You may connect to GPU by checking `Runtime` > `Change runtime type` > `Hardware accelerator: GPU`\n",
+        "* Next, install the latest PyPi distribution of snnTorch by clicking into the following cell and pressing `Shift+Enter`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "BgBRVUtpRp74"
+      },
+      "source": [
+        "!pip install snntorch"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Zm-D2lthRp75"
+      },
+      "source": [
+        "## 1. Setting up the Static MNIST Dataset\n",
+        "### 1.1. Import packages and setup environment"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "sEygpdc8Rp76"
+      },
+      "source": [
+        "import snntorch as snn\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torchvision import datasets, transforms\n",
+        "import numpy as np\n",
+        "import itertools\n",
+        "import matplotlib.pyplot as plt"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "RcX7J9vVRp76"
+      },
+      "source": [
+        "### 1.2 Define network and SNN parameters\n",
+        "We will use a 784-1000-10 FCN architecture for a sequence of 25 time steps.\n",
+        "\n",
+        "* `alpha` is the decay rate of the synaptic current of a neuron\n",
+        "* `beta` is the decay rate of the membrane potential of a neuron"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "bKEj2hucRp77"
+      },
+      "source": [
+        "# Network Architecture\n",
+        "num_inputs = 28*28\n",
+        "num_hidden = 1000\n",
+        "num_outputs = 10\n",
+        "\n",
+        "# Training Parameters\n",
+        "batch_size=128\n",
+        "data_path='/data/mnist'\n",
+        "\n",
+        "# Temporal Dynamics\n",
+        "num_steps = 25\n",
+        "time_step = 1e-3\n",
+        "tau_mem = 3e-3\n",
+        "tau_syn = 2.2e-3\n",
+        "alpha = float(np.exp(-time_step/tau_syn))\n",
+        "beta = float(np.exp(-time_step/tau_mem))\n",
+        "\n",
+        "dtype = torch.float\n",
+        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "gM_hcwDIRp78"
+      },
+      "source": [
+        "### 1.3 Download MNIST Dataset\n",
+        "To see how to construct a validation set, refer to Tutorial 1."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "0xwYb15xRp79"
+      },
+      "source": [
+        "# Define a transform\n",
+        "transform = transforms.Compose([\n",
+        "            transforms.Resize((28, 28)),\n",
+        "            transforms.Grayscale(),\n",
+        "            transforms.ToTensor(),\n",
+        "            transforms.Normalize((0,), (1,))])\n",
+        "\n",
+        "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n",
+        "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "DIg3vLmURp7-"
+      },
+      "source": [
+        "### 1.4 Create DataLoaders"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "K4DF-odMRp7_"
+      },
+      "source": [
+        "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
+        "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "pD4Dw-RoRp7_"
+      },
+      "source": [
+        "## 2. Define Network\n",
+        "snnTorch treats neurons as activations with recurrent connections. This allows for smooth integration with PyTorch.\n",
+        "There are a few useful neuron models and surrogate gradient functions which approximate the gradient of spikes.\n",
+        "\n",
+        "Our network will use one type of neuron model and one surrogate gradient:\n",
+        "1. `snntorch.Stein` is a basic leaky integrate and fire (LIF) neuron. Specifically, it assumes instantaneous rise times for synaptic current and membrane potential.\n",
+        "2. `snntorch.FastSigmoidSurrogate` defines separate forward and backward functions. The forward function is a Heaviside step function for spike generation. The backward function is the derivative of a fast sigmoid function, to ensure continuous differentiability.\n",
+        "The `FastSigmoidSurrogate` function has been adapted from:\n",
+        "\n",
+        ">Neftci, E. O., Mostafa, H., and Zenke, F. (2019) Surrogate Gradient Learning in Spiking Neural Networks. https://arxiv.org/abs/1901/09948"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "Sf9RdE9jRp8A"
+      },
+      "source": [
+        "# from snntorch import surrogate\n",
+        "#\n",
+        "# spike_grad = surrogate.FastSigmoid.apply\n",
+        "# snn.slope = 50"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Og9e57W0Rp8B"
+      },
+      "source": [
+        "The surrogate is passed to `spike_grad` and overrides the default gradient of the Heaviside step function.\n",
+        "If we did not override the default gradient, (zero everywhere, except for $x=1$ where it is technically infinite but clipped to 1 here), then learning would not take place for as long as the neuron was not emitting post-synaptic spikes.\n",
+        "\n",
+        "`snn.slope` defines the slope of the backward surrogate.\n",
+        "\n",
+        "TO-DO: Include visualisation."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Yo-x48mARp8C"
+      },
+      "source": [
+        "Now we can define our spiking neural network (SNN).\n",
+        "Creating an instance of the `Stein` neuron requires two arguments and two optional arguments:\n",
+        "1. $I_{syn}$ decay rate, $\\alpha$,\n",
+        "2. $V_{mem}$ decay rate, $\\beta$,\n",
+        "3. the surrogate spiking function, `spike_grad` (*default*: the gradient of the Heaviside function), and\n",
+        "4. the threshold for spiking, (*default*: 1.0).\n",
+        "\n",
+        "snnTorch treats the LIF neuron as a recurrent activation. Therefore, it requires initialization of its internal states.\n",
+        "For each layer, we initialize the synaptic current `syn1` and `syn2`, the membrane potential `mem1` and `mem2`, and the post-synaptic spikes `spk1` and `spk2` to zero.\n",
+        "A class method `init_stein` will take care of this.\n",
+        "\n",
+        "For rate coding, the final layer of spikes and membrane potential are used to determine accuracy and loss, respectively.\n",
+        "So their historical values are recorded in `spk2_rec` and `mem2_rec`.\n",
+        "\n",
+        "Keep in mind, the dataset we are using is just static MNIST. I.e., it is *not* time-varying.\n",
+        "Therefore, we pass the same MNIST sample to the input at each time step.\n",
+        "This is handled in the line `cur1 = self.fc1(x)`, where `x` is the same input over the whole for-loop."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "P6RHCnXMRp8D"
+      },
+      "source": [
+        "# Define Network\n",
+        "class Net(nn.Module):\n",
+        "    def __init__(self):\n",
+        "        super().__init__()\n",
+        "\n",
+        "    # initialize layers\n",
+        "        self.fc1 = nn.Linear(num_inputs, num_hidden)\n",
+        "        self.lif1 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "        self.fc2 = nn.Linear(num_hidden, num_outputs)\n",
+        "        self.lif2 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        spk1, syn1, mem1 = self.lif1.init_stein(batch_size, num_hidden)\n",
+        "        spk2, syn2, mem2 = self.lif2.init_stein(batch_size, num_outputs)\n",
+        "\n",
+        "        spk2_rec = []\n",
+        "        mem2_rec = []\n",
+        "\n",
+        "        for step in range(num_steps):\n",
+        "            cur1 = self.fc1(x)\n",
+        "            spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)\n",
+        "            cur2 = self.fc2(spk1)\n",
+        "            spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)\n",
+        "\n",
+        "            spk2_rec.append(spk2)\n",
+        "            mem2_rec.append(mem2)\n",
+        "\n",
+        "        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n",
+        "\n",
+        "net = Net().to(device)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "386KNHG7Rp8E"
+      },
+      "source": [
+        "## 3. Training\n",
+        "Time for training! Let's first define a couple of functions to print out test/train accuracy."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "cOKKbUnDRp8F"
+      },
+      "source": [
+        "def print_batch_accuracy(data, targets, train=False):\n",
+        "    output, _ = net(data.view(batch_size, -1))\n",
+        "    _, idx = output.sum(dim=0).max(1)\n",
+        "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
+        "\n",
+        "    if train:\n",
+        "        print(f\"Train Set Accuracy: {acc}\")\n",
+        "    else:\n",
+        "        print(f\"Test Set Accuracy: {acc}\")\n",
+        "\n",
+        "def train_printer():\n",
+        "    print(f\"Epoch {epoch}, Minibatch {minibatch_counter}\")\n",
+        "    print(f\"Train Set Loss: {loss_hist[counter]}\")\n",
+        "    print(f\"Test Set Loss: {test_loss_hist[counter]}\")\n",
+        "    print_batch_accuracy(data_it, targets_it, train=True)\n",
+        "    print_batch_accuracy(testdata_it, testtargets_it, train=False)\n",
+        "    print(\"\\n\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "mZqfCe0KRp8J"
+      },
+      "source": [
+        "### 3.1 Optimizer & Loss\n",
+        "* *Output Activation*: We'll apply the softmax function to the membrane potentials of the output layer, rather than the spikes.\n",
+        "* *Loss*: This will then be used to calculate the negative log-likelihood loss.\n",
+        "By encouraging the membrane of the correct neuron class to reach the threshold, we expect that neuron will fire more frequently.\n",
+        "The loss could be applied to the spike count as well, but the membrane is  continuous whereas spike count is discrete.\n",
+        "* *Optimizer*: The Adam optimizer is used for weight updates.\n",
+        "* *Accuracy*: Accuracy is measured by counting the spikes of the output neurons. The neuron that fires the most frequently will be our predicted class."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "UMZ-4uGlRp8K"
+      },
+      "source": [
+        "optimizer = torch.optim.Adam(net.parameters(), lr=2e-4, betas=(0.9, 0.999))\n",
+        "log_softmax_fn = nn.LogSoftmax(dim=-1)\n",
+        "loss_fn = nn.NLLLoss()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "UerqLkd1Rp8K"
+      },
+      "source": [
+        "### 3.2 Training Loop\n",
+        "Now just sit back, relax, and wait for convergence."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "76CEHI2xRp8L"
+      },
+      "source": [
+        "loss_hist = []\n",
+        "test_loss_hist = []\n",
+        "counter = 0\n",
+        "\n",
+        "# Outer training loop\n",
+        "for epoch in range(3):\n",
+        "    minibatch_counter = 0\n",
+        "    train_batch = iter(train_loader)\n",
+        "\n",
+        "    # Minibatch training loop\n",
+        "    for data_it, targets_it in train_batch:\n",
+        "        data_it = data_it.to(device)\n",
+        "        targets_it = targets_it.to(device)\n",
+        "\n",
+        "        output, mem_rec = net(data_it.view(batch_size, -1))\n",
+        "        log_p_y = log_softmax_fn(mem_rec)\n",
+        "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
+        "\n",
+        "        # Sum loss over time steps: BPTT\n",
+        "        for step in range(num_steps):\n",
+        "          loss_val += loss_fn(log_p_y[step], targets_it)\n",
+        "\n",
+        "        # Gradient calculation\n",
+        "        optimizer.zero_grad()\n",
+        "        loss_val.backward(retain_graph=True)\n",
+        "\n",
+        "        # Weight Update\n",
+        "        nn.utils.clip_grad_norm_(net.parameters(), 1) # gradient clipping\n",
+        "        optimizer.step()\n",
+        "\n",
+        "        # Store loss history for future plotting\n",
+        "        loss_hist.append(loss_val.item())\n",
+        "\n",
+        "        # Test set\n",
+        "        test_data = itertools.cycle(test_loader)\n",
+        "        testdata_it, testtargets_it = next(test_data)\n",
+        "        testdata_it = testdata_it.to(device)\n",
+        "        testtargets_it = testtargets_it.to(device)\n",
+        "\n",
+        "        # Test set forward pass\n",
+        "        test_output, test_mem_rec = net(testdata_it.view(batch_size, -1))\n",
+        "\n",
+        "        # Test set loss\n",
+        "        log_p_ytest = log_softmax_fn(test_mem_rec)\n",
+        "        log_p_ytest = log_p_ytest.sum(dim=0)\n",
+        "        loss_val_test = loss_fn(log_p_ytest, testtargets_it)\n",
+        "        test_loss_hist.append(loss_val_test.item())\n",
+        "\n",
+        "        # Print test/train loss/accuracy\n",
+        "        if counter % 50 == 0:\n",
+        "            train_printer()\n",
+        "        minibatch_counter += 1\n",
+        "        counter += 1\n",
+        "\n",
+        "loss_hist_true_grad = loss_hist\n",
+        "test_loss_hist_true_grad = test_loss_hist"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "fdOuwCWHRp8L"
+      },
+      "source": [
+        "## 4. Results\n",
+        "### 4.1 Plot Training/Test Loss"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "WJGSBq6zRp8M"
+      },
+      "source": [
+        "# Plot Loss\n",
+        "fig = plt.figure(facecolor=\"w\", figsize=(10, 5))\n",
+        "plt.plot(loss_hist)\n",
+        "plt.plot(test_loss_hist)\n",
+        "plt.legend([\"Test Loss\", \"Train Loss\"])\n",
+        "plt.xlabel(\"Epoch\")\n",
+        "plt.ylabel(\"Loss\")\n",
+        "plt.show()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "h7xb37iHRp8N"
+      },
+      "source": [
+        "### 4.2 Test Set Accuracy\n",
+        "This function just iterates over all minibatches to obtain a measure of accuracy over the full 10,000 samples in the test set."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "R1ReGuNURp8N"
+      },
+      "source": [
+        "total = 0\n",
+        "correct = 0\n",
+        "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "  net.eval()\n",
+        "  for data in test_loader:\n",
+        "    images, labels = data\n",
+        "    images = images.to(device)\n",
+        "    labels = labels.to(device)\n",
+        "\n",
+        "    # If current batch matches batch_size, just do the usual thing\n",
+        "    if images.size()[0] == batch_size:\n",
+        "      outputs, _ = net(images.view(batch_size, -1))\n",
+        "\n",
+        "    # If current batch does not match batch_size (e.g., is the final minibatch),\n",
+        "    # modify batch_size in a temp variable and restore it at the end of the else block\n",
+        "    else:\n",
+        "      temp_bs = batch_size\n",
+        "      batch_size = images.size()[0]\n",
+        "      outputs, _ = net(images.view(images.size()[0], -1))\n",
+        "      batch_size = temp_bs\n",
+        "\n",
+        "    _, predicted = outputs.sum(dim=0).max(1)\n",
+        "    total += labels.size(0)\n",
+        "    correct += (predicted == labels).sum().item()\n",
+        "\n",
+        "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
+        "print(f\"Test Set Accuracy: {100 * correct / total}%\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "zNJjSKATRp8N"
+      },
+      "source": [
+        "Voila! That's it for static MNIST."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "2qUxjHwBRp8f"
+      },
+      "source": [
+        "## 5. Spiking MNIST\n",
+        "Part of the appeal of SNNs is their ability to handle time-varying spiking data. So let's use rate-coding to convert MNIST into spiking MNIST using the `spikegen` module in the previous tutorial, and train our network with that instead."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "8K7_C-2rRp8g"
+      },
+      "source": [
+        "from snntorch import spikegen\n",
+        "\n",
+        "# MNIST to spiking-MNIST\n",
+        "spike_data, spike_targets = spikegen.rate(data_it, targets_it, num_outputs=num_outputs, num_steps=num_steps, gain=1,\n",
+        "                                          offset=0, convert_targets=False, temporal_targets=False)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "BvaNfns_Rp8h"
+      },
+      "source": [
+        "### 5.1 Visualiser\n",
+        "Just so you're damn sure it's a spiking input."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "ll2D4jtdRp8i"
+      },
+      "source": [
+        "!pip install celluloid # matplotlib animations made easy"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "cJQ0XUAZRp8i"
+      },
+      "source": [
+        "from celluloid import Camera\n",
+        "from IPython.display import HTML\n",
+        "\n",
+        "# Animator\n",
+        "spike_data_sample = spike_data[:, 0, 0].cpu()\n",
+        "\n",
+        "fig, ax = plt.subplots()\n",
+        "camera = Camera(fig)\n",
+        "plt.axis('off')\n",
+        "\n",
+        "for step in range(num_steps):\n",
+        "    im = ax.imshow(spike_data_sample[step, :, :], cmap='plasma')\n",
+        "    camera.snap()\n",
+        "\n",
+        "# interval=40 specifies 40ms delay between frames\n",
+        "a = camera.animate(interval=40)\n",
+        "HTML(a.to_html5_video())"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "T-dMHLBbRp8k"
+      },
+      "source": [
+        "print(spike_targets[0])"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "I63CROTaRp8k"
+      },
+      "source": [
+        "## 6. Define Network\n",
+        "The network is the same as before. The one difference is that the for-loop iterates through the first dimension of the input:\n",
+        "`cur1 = self.fc1(x[step])`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "a_elvRYIRp8k"
+      },
+      "source": [
+        "spike_grad = snn.FastSigmoidSurrogate.apply\n",
+        "snn.slope = 50 # The lower the slope, the smoother the gradient\n",
+        "\n",
+        "# Define Network\n",
+        "class Net(nn.Module):\n",
+        "    def __init__(self):\n",
+        "        super().__init__()\n",
+        "\n",
+        "        # Initialize layers\n",
+        "        self.fc1 = nn.Linear(num_inputs, num_hidden)\n",
+        "        self.lif1 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "        self.fc2 = nn.Linear(num_hidden, num_outputs)\n",
+        "        self.lif2 = snn.Stein(alpha=alpha, beta=beta, spike_grad=spike_grad)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        # Initialize hidden states + output spike at t=0\n",
+        "        spk1, syn1, mem1 = self.lif1.init_stein(batch_size, num_hidden)\n",
+        "        spk2, syn2, mem2 = self.lif2.init_stein(batch_size, num_outputs)\n",
+        "\n",
+        "        spk2_rec = []\n",
+        "        mem2_rec = []\n",
+        "\n",
+        "        for step in range(num_steps):\n",
+        "            cur1 = self.fc1(x[step])\n",
+        "            spk1, syn1, mem1 = self.lif1(cur1, syn1, mem1)\n",
+        "            cur2 = self.fc2(spk1)\n",
+        "            spk2, syn2, mem2 = self.lif2(cur2, syn2, mem2)\n",
+        "\n",
+        "            spk2_rec.append(spk2)\n",
+        "            mem2_rec.append(mem2)\n",
+        "\n",
+        "        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n",
+        "\n",
+        "net = Net().to(device)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "E8utR7I9Rp8l"
+      },
+      "source": [
+        "## 7. Training\n",
+        "We make a slight modification to our print-out functions to handle the new first dimension of the input:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "FlkQt3skRp8l"
+      },
+      "source": [
+        "def print_batch_accuracy(data, targets, train=False):\n",
+        "    output, _ = net(data.view(num_steps, batch_size, -1))\n",
+        "    _, idx = output.sum(dim=0).max(1)\n",
+        "    acc = np.mean((targets == idx).detach().cpu().numpy())\n",
+        "\n",
+        "    if train:\n",
+        "        print(f\"Train Set Accuracy: {acc}\")\n",
+        "    else:\n",
+        "        print(f\"Test Set Accuracy: {acc}\")\n",
+        "\n",
+        "def train_printer():\n",
+        "    print(f\"Epoch {epoch}, Minibatch {minibatch_counter}\")\n",
+        "    print(f\"Train Set Loss: {loss_hist[counter]}\")\n",
+        "    print(f\"Test Set Loss: {test_loss_hist[counter]}\")\n",
+        "    print_batch_accuracy(spike_data, spike_targets, train=True)\n",
+        "    print_batch_accuracy(test_spike_data, test_spike_targets, train=False)\n",
+        "    print(\"\\n\")\n"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "vbve35sDRp8l"
+      },
+      "source": [
+        "### 7.1 Optimizer & Loss\n",
+        "We'll keep our optimizer and loss the exact same as the static MNIST case."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "10CH8RKHRp8l"
+      },
+      "source": [
+        "optimizer = torch.optim.Adam(net.parameters(), lr=2e-4, betas=(0.9, 0.999))\n",
+        "log_softmax_fn = nn.LogSoftmax(dim=-1)\n",
+        "loss_fn = nn.NLLLoss()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "pycharm": {
+          "name": "#%% md\n"
+        },
+        "id": "Ef0TM6yIRp8m"
+      },
+      "source": [
+        "### 7.2 Training Loop\n",
+        "The training loop is identical to the static MNIST case, but we pass each minibatch through `spikegen.rate` before running it through the feedforward network."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "ISLzPpwvRp8m"
+      },
+      "source": [
+        "loss_hist = []\n",
+        "test_loss_hist = []\n",
+        "counter = 0\n",
+        "\n",
+        "# Outer training loop\n",
+        "for epoch in range(3):\n",
+        "    minibatch_counter = 0\n",
+        "    data = iter(train_loader)\n",
+        "\n",
+        "    # Minibatch training loop\n",
+        "    for data_it, targets_it in data:\n",
+        "        data_it = data_it.to(device)\n",
+        "        targets_it = targets_it.to(device)\n",
+        "\n",
+        "        # Spike generator\n",
+        "        spike_data, spike_targets = spikegen.rate(data_it, targets_it, num_outputs=num_outputs, num_steps=num_steps,\n",
+        "                                                  gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "        # Forward pass\n",
+        "        output, mem_rec = net(spike_data.view(num_steps, batch_size, -1))\n",
+        "        log_p_y = log_softmax_fn(mem_rec)\n",
+        "        loss_val = torch.zeros((1), dtype=dtype, device=device)\n",
+        "\n",
+        "        # Sum loss over time steps to perform BPTT\n",
+        "        for step in range(num_steps):\n",
+        "          loss_val += loss_fn(log_p_y[step], targets_it)\n",
+        "\n",
+        "        # Gradient Calculation\n",
+        "        optimizer.zero_grad()\n",
+        "        loss_val.backward(retain_graph=True)\n",
+        "        nn.utils.clip_grad_norm_(net.parameters(), 1)\n",
+        "\n",
+        "        # Weight Update\n",
+        "        optimizer.step()\n",
+        "\n",
+        "        # Store Loss history\n",
+        "        loss_hist.append(loss_val.item())\n",
+        "\n",
+        "        # Test set\n",
+        "        test_data = itertools.cycle(test_loader)\n",
+        "        testdata_it, testtargets_it = next(test_data)\n",
+        "        testdata_it = testdata_it.to(device)\n",
+        "        testtargets_it = testtargets_it.to(device)\n",
+        "\n",
+        "        # Test set spike conversion\n",
+        "        test_spike_data, test_spike_targets = spikegen.rate(testdata_it, testtargets_it, num_outputs=num_outputs,\n",
+        "                                                            num_steps=num_steps, gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "        # Test set forward pass\n",
+        "        test_output, test_mem_rec = net(test_spike_data.view(num_steps, batch_size, -1))\n",
+        "\n",
+        "        # Test set loss\n",
+        "        log_p_ytest = log_softmax_fn(test_mem_rec)\n",
+        "        log_p_ytest = log_p_ytest.sum(dim=0)\n",
+        "        loss_val_test = loss_fn(log_p_ytest, test_spike_targets)\n",
+        "        test_loss_hist.append(loss_val_test.item())\n",
+        "\n",
+        "        # Print test/train loss/accuracy\n",
+        "        if counter % 50 == 0:\n",
+        "            train_printer()\n",
+        "        minibatch_counter += 1\n",
+        "        counter += 1\n",
+        "\n",
+        "loss_hist_true_grad = loss_hist\n",
+        "test_loss_hist_true_grad = test_loss_hist"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "EuRL61AnRp8m"
+      },
+      "source": [
+        "## 8. Spiking MNIST Results\n",
+        "### 8.1 Plot Training/Test Loss"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "fl5griHBRp8n"
+      },
+      "source": [
+        "# Plot Loss\n",
+        "fig = plt.figure(facecolor=\"w\", figsize=(10, 5))\n",
+        "plt.plot(loss_hist)\n",
+        "plt.plot(test_loss_hist)\n",
+        "plt.legend([\"Test Loss\", \"Train Loss\"])\n",
+        "plt.xlabel(\"Epoch\")\n",
+        "plt.ylabel(\"Loss\")\n",
+        "plt.show()"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "IZMH5UOlRp8n"
+      },
+      "source": [
+        "### 8.2 Test Set Accuracy"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "pycharm": {
+          "name": "#%%\n"
+        },
+        "id": "eO2Krz1mRp8n"
+      },
+      "source": [
+        "total = 0\n",
+        "correct = 0\n",
+        "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "  net.eval()\n",
+        "  for data in test_loader:\n",
+        "    images, labels = data\n",
+        "    images = images.to(device)\n",
+        "    labels = labels.to(device)\n",
+        "\n",
+        "    # If current batch matches batch_size, just do the usual thing\n",
+        "    if images.size()[0] == batch_size:\n",
+        "      spike_test, spike_targets = spikegen.rate(images, labels, num_outputs=num_outputs, num_steps=num_steps,\n",
+        "                                                            gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "      outputs, _ = net(spike_test.view(num_steps, batch_size, -1))\n",
+        "\n",
+        "    # If current batch does not match batch_size (e.g., is the final minibatch),\n",
+        "    # modify batch_size in a temp variable and restore it at the end of the else block\n",
+        "    else:\n",
+        "      temp_bs = batch_size\n",
+        "      batch_size = images.size()[0]\n",
+        "      spike_test, spike_targets = spikegen.rate(images, labels, num_outputs=num_outputs, num_steps=num_steps,\n",
+        "                                                gain=1, offset=0, convert_targets=False, temporal_targets=False)\n",
+        "\n",
+        "      outputs, _ = net(spike_test.view(num_steps, images.size()[0], -1))\n",
+        "      batch_size = temp_bs\n",
+        "\n",
+        "    _, predicted = outputs.sum(dim=0).max(1)\n",
+        "    total += spike_targets.size(0)\n",
+        "    correct += (predicted == spike_targets).sum().item()\n",
+        "\n",
+        "print(f\"Total correctly classified test set images: {correct}/{total}\")\n",
+        "print(f\"Test Set Accuracy: {100 * correct / total}%\")"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "collapsed": false,
+        "id": "ztM4ogAqRp8n"
+      },
+      "source": [
+        "That's all for now!\n",
+        "Next time, we'll introduce how to use spiking convolutional layers to improve accuracy."
+      ]
+    }
+  ]
+}
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
index e68840a0..1a11855c 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.1.8
+current_version = 0.1.15
 commit = True
 tag = True
 
diff --git a/setup.py b/setup.py
index 0bdd6f63..ccb931d1 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
     history = history_file.read()
 
 # fmt: off
-__version__ = '0.1.8'
+__version__ = '0.1.15'
 # fmt: on
 
 requirements = [
diff --git a/snntorch/__init__.py b/snntorch/__init__.py
index 7928e7eb..7000eb1b 100644
--- a/snntorch/__init__.py
+++ b/snntorch/__init__.py
@@ -15,7 +15,15 @@ class LIF(nn.Module):
     """Each :mod:`snntorch.LIF` neuron (e.g., :mod:`snntorch.Stein`) will populate the :mod:`snntorch.LIF.instances` list with a new entry.
     The list is used to initialize and clear neuron states when the argument `init_hidden=True`."""
 
-    def __init__(self, alpha, beta, threshold=1.0, spike_grad=None, inhibition=False):
+    def __init__(
+        self,
+        alpha,
+        beta,
+        threshold=1.0,
+        spike_grad=None,
+        inhibition=False,
+        reset_mechanism="subtract",
+    ):
         super(LIF, self).__init__()
         LIF.instances.append(self)
 
@@ -23,18 +31,25 @@ def __init__(self, alpha, beta, threshold=1.0, spike_grad=None, inhibition=False
         self.beta = beta
         self.threshold = threshold
         self.inhibition = inhibition
+        self.reset_mechanism = reset_mechanism
 
         if spike_grad is None:
             self.spike_grad = self.Heaviside.apply
         else:
             self.spike_grad = spike_grad
 
+        if reset_mechanism != "subtract" and reset_mechanism != "zero":
+            raise ValueError(
+                "reset_mechanism must be set to either 'subtract' or 'zero'."
+            )
+
     def fire(self, mem):
         """Generates spike if mem > threshold.
         Returns spk and reset."""
         mem_shift = mem - self.threshold
         spk = self.spike_grad(mem_shift).to(device)
         reset = spk.clone().detach()
+
         return spk, reset
 
     def fire_inhibition(self, batch_size, mem):
@@ -47,8 +62,8 @@ def fire_inhibition(self, batch_size, mem):
         mask_spk1 = torch.zeros_like(spk_tmp)
         mask_spk1[torch.arange(batch_size), index] = 1
         spk = spk_tmp * mask_spk1.to(device)
-
         reset = spk.clone().detach()
+
         return spk, reset
 
     @classmethod
@@ -118,8 +133,7 @@ class Heaviside(torch.autograd.Function):
         @staticmethod
         def forward(ctx, input_):
             ctx.save_for_backward(input_)
-            out = torch.zeros_like(input_)
-            out[input_ >= 0] = 1.0
+            out = (input_ > 0).float()
             return out
 
         @staticmethod
@@ -141,15 +155,23 @@ class Stein(LIF):
     Synaptic current and membrane potential decay exponentially with rates of alpha and beta, respectively.
     For :math:`U[T] > U_{\\rm thr} ⇒ S[T+1] = 1`.
 
+     If `reset_mechanism = "subtract"`, then :math:`U[t+1]` will have `threshold` subtracted from it whenever the neuron emits a spike:
     .. math::
 
             I_{\\rm syn}[t+1] = αI_{\\rm syn}[t] + I_{\\rm in}[t+1] \\\\
-            U[t+1] = βU[t] + I_{\\rm syn}[t+1] - R
+            U[t+1] = βU[t] + I_{\\rm syn}[t+1] - RU_{\\rm thr}
+
+    If `reset_mechanism = "zero"`, then :math:`U[t+1]` will be set to `0` whenever the neuron emits a spike:
+    .. math::
+
+            I_{\\rm syn}[t+1] = αI_{\\rm syn}[t] + I_{\\rm in}[t+1] \\\\
+            U[t+1] = βU[t] + I_{\\rm syn}[t+1] - R(βU[t] + I_{\\rm syn}[t+1])
 
     * :math:`I_{\\rm syn}` - Synaptic current
     * :math:`I_{\\rm in}` - Input current
     * :math:`U` - Membrane potential
-    * :math:`R` - Reset mechanism
+    * :math:`U_{\\rm thr}` - Membrane threshold
+    * :math:`R` - Reset mechanism: if active, :math:`R = U_{\\rm thr}`, otherwise :math:`R = 0`
     * :math:`α` - Synaptic current decay rate
     * :math:`β` - Membrane potential decay rate
 
@@ -197,8 +219,11 @@ def __init__(
         batch_size=False,
         hidden_init=False,
         inhibition=False,
+        reset_mechanism="subtract",
     ):
-        super(Stein, self).__init__(alpha, beta, threshold, spike_grad, inhibition)
+        super(Stein, self).__init__(
+            alpha, beta, threshold, spike_grad, inhibition, reset_mechanism
+        )
 
         self.num_inputs = num_inputs
         self.batch_size = batch_size
@@ -234,7 +259,12 @@ def forward(self, input_, syn, mem):
             else:
                 spk, reset = self.fire(mem)
             syn = self.alpha * syn + input_
-            mem = self.beta * mem + syn - reset
+
+            if self.reset_mechanism == "subtract":
+                mem = self.beta * mem + syn - reset * self.threshold
+
+            elif self.reset_mechanism == "zero":
+                mem = self.beta * mem + syn - reset * (self.beta * mem + syn)
 
             return spk, syn, mem
 
@@ -245,7 +275,16 @@ def forward(self, input_, syn, mem):
             else:
                 self.spk, self.reset = self.fire(self.mem)
             self.syn = self.alpha * self.syn + input_
-            self.mem = self.beta * self.mem + self.syn - self.reset
+
+            if self.reset_mechanism == "subtract":
+                self.mem = self.beta * self.mem + self.syn - self.reset * self.threshold
+
+            elif self.reset_mechanism == "zero":
+                self.mem = (
+                    self.beta * self.mem
+                    + self.syn
+                    - self.reset * (self.beta * self.mem + self.syn)
+                )
 
             return self.spk, self.syn, self.mem
 
@@ -281,17 +320,28 @@ class SRM0(LIF):
 
     .. warning:: For a positive input current to induce a positive membrane response, ensure :math:`α > β`.
 
+    If `reset_mechanism = "subtract"`, then :math:`I_{\\rm syn-pre}, I_{\\rm syn-post}` will both have `threshold` subtracted from them whenever the neuron emits a spike:
     .. math::
 
-            I_{\\rm syn-pre}[t+1] = (αI_{\\rm syn-pre}[t] + I_{\\rm in}[t+1])(1-R) \\\\
-            I_{\\rm syn-post}[t+1] = (βI_{\\rm syn-post}[t] - I_{\\rm in}[t+1])(1-R) \\\\
-            U[t+1] = τ_{\\rm SRM}(I_{\\rm syn-pre}[t+1] + I_{\\rm syn-post}[t+1])(1-R)
+            I_{\\rm syn-pre}[t+1] = (αI_{\\rm syn-pre}[t] + I_{\\rm in}[t+1]) - R(αI_{\\rm syn-pre}[t] + I_{\\rm in}[t+1]) \\\\
+            I_{\\rm syn-post}[t+1] = (βI_{\\rm syn-post}[t] - I_{\\rm in}[t+1]) - R(βI_{\\rm syn-post}[t] - I_{\\rm in}[t+1]) \\\\
+            U_{\\rm residual}[t+1] = R(U[t]U_{\\rm thr}) + U[t]/τ_{\\rm SRM} \\\\
+            U[t+1] = τ_{\\rm SRM}(I_{\\rm syn-pre}[t+1] + I_{\\rm syn-post}[t+1]) + U_{\\rm residual}[t+1]
+
+    If `reset_mechanism = "zero"`, then :math:`I_{\\rm syn-pre}, I_{\\rm syn-post}` will both be set to `0` whenever the neuron emits a spike:
+    .. math::
+
+            I_{\\rm syn-pre}[t+1] = (αI_{\\rm syn-pre}[t] + I_{\\rm in}[t+1]) - R(αI_{\\rm syn-pre}[t] + I_{\\rm in}[t+1]) \\\\
+            I_{\\rm syn-post}[t+1] = (βI_{\\rm syn-post}[t] - I_{\\rm in}[t+1]) - R(βI_{\\rm syn-post}[t] - I_{\\rm in}[t+1]) \\\\
+            U[t+1] = τ_{\\rm SRM}(I_{\\rm syn-pre}[t+1] + I_{\\rm syn-post}[t+1])
 
     * :math:`I_{\\rm syn-pre}` - Pre-synaptic current
     * :math:`I_{\\rm syn-post}` - Post-synaptic current
     * :math:`I_{\\rm in}` - Input current
     * :math:`U` - Membrane potential
-    * :math:`R` - Reset mechanism
+    * :math:`U_{\\rm thr}` - Membrane threshold
+    * :math:`U_{\\rm residual}` - Residual membrane potential after reset by subtraction
+    * :math:`R` - Reset mechanism, :math:`R = 1` if spike occurs, otherwise :math:`R = 0`
     * :math:`α` - Pre-synaptic current decay rate
     * :math:`β` - Post-synaptic current decay rate
     * :math:`τ_{\\rm SRM} = \\frac{log(α)}{log(β)} - log(α) + 1`
@@ -338,8 +388,11 @@ def __init__(
         batch_size=False,
         hidden_init=False,
         inhibition=False,
+        reset_mechanism="subtract",
     ):
-        super(SRM0, self).__init__(alpha, beta, threshold, spike_grad, inhibition)
+        super(SRM0, self).__init__(
+            alpha, beta, threshold, spike_grad, inhibition, reset_mechanism
+        )
 
         self.num_inputs = num_inputs
         self.batch_size = batch_size
@@ -376,33 +429,92 @@ def __init__(
                 "beta cannot be '1' otherwise ZeroDivisionError occurs: tau_srm = log(alpha)/log(beta) - log(alpha) + 1"
             )
 
+        if reset_mechanism == "subtract":
+            self.mem_residual = False
+
         self.tau_srm = np.log(self.alpha) / (np.log(self.beta) - np.log(self.alpha)) + 1
 
     def forward(self, input_, syn_pre, syn_post, mem):
+
         # if hidden states are passed externally
         if not self.hidden_init:
+
             if self.inhibition:
                 spk, reset = self.fire_inhibition(self.batch_size, mem)
+
             else:
                 spk, reset = self.fire(mem)
-            syn_pre = (self.alpha * syn_pre + input_) * (1 - reset)
-            syn_post = (self.beta * syn_post - input_) * (1 - reset)
-            mem = self.tau_srm * (syn_pre + syn_post) * (1 - reset) + (
-                mem * reset - reset
-            )
+
+            # if neuron fires, subtract threhsold from neuron
+            if self.reset_mechanism == "subtract":
+
+                if self.mem_residual is False:
+                    self.mem_residual = torch.zeros_like(mem)
+
+                syn_pre = (self.alpha * syn_pre + input_) - reset * (
+                    self.alpha * syn_pre + input_
+                )
+                syn_post = (self.beta * syn_post - input_) - reset * (
+                    self.beta * syn_post - input_
+                )
+                # The residual of (mem - threshold) decays separately
+                self.mem_residual = reset * (mem - self.threshold) + (
+                    self.mem_residual / self.tau_srm
+                )
+                mem = self.tau_srm * (syn_pre + syn_post) + self.mem_residual
+
+            # if neuron fires, reset membrane to zero
+            elif self.reset_mechanism == "zero":
+                syn_pre = (self.alpha * syn_pre + input_) - reset * (
+                    self.alpha * syn_pre + input_
+                )
+                syn_post = (self.beta * syn_post - input_) - reset * (
+                    self.beta * syn_post - input_
+                )
+                mem = self.tau_srm * (syn_pre + syn_post)
+
             return spk, syn_pre, syn_post, mem
 
         # if hidden states and outputs are instance variables
         if self.hidden_init:
+
             if self.inhibition:
                 self.spk, self.reset = self.fire_inhibition(self.batch_size, self.mem)
+
             else:
                 self.spk, self.reset = self.fire(self.mem)
-            self.syn_pre = (self.alpha * self.syn_pre + input_) * (1 - self.reset)
-            self.syn_post = (self.beta * self.syn_post - input_) * (1 - self.reset)
-            self.mem = self.tau_srm * (self.syn_pre + self.syn_post) * (
-                1 - self.reset
-            ) + (self.mem * self.reset - self.reset)
+
+            # if neuron fires, subtract threhsold from neuron
+            if self.reset_mechanism == "subtract":
+
+                if self.mem_residual is False:
+                    self.mem_residual = torch.zeros_like(mem)
+
+                self.syn_pre = (self.alpha * self.syn_pre + input_) - self.reset * (
+                    self.alpha * self.syn_pre + input_
+                )
+                syn_post = (self.beta * self.syn_post - input_) - self.reset * (
+                    self.beta * self.syn_post - input_
+                )
+                # The residual of (mem - threshold) decays separately
+                self.mem_residual = self.reset * (self.mem - self.threshold) + (
+                    self.mem_residual / self.tau_srm
+                )
+                self.mem = (
+                    self.tau_srm * (self.syn_pre + self.syn_post) + self.mem_residual
+                )
+
+            # if neuron fires, reset membrane to zero
+            elif self.reset_mechanism == "zero":
+
+                syn_pre = (self.alpha * syn_pre + input_) - self.reset * (
+                    self.alpha * syn_pre + input_
+                )
+                syn_post = (self.beta * syn_post - input_) - self.reset * (
+                    self.beta * syn_post - input_
+                )
+                self.mem = self.tau_srm * (syn_pre + syn_post)
+
             return self.spk, self.syn_pre, self.syn_post, self.mem
 
     # cool forward function that resulted in burst firing - worth exploring
diff --git a/snntorch/surrogate.py b/snntorch/surrogate.py
index 10686cb5..bc7acc5e 100644
--- a/snntorch/surrogate.py
+++ b/snntorch/surrogate.py
@@ -33,8 +33,7 @@ class FastSigmoid(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input_):
         ctx.save_for_backward(input_)
-        out = torch.zeros_like(input_)
-        out[input_ > 0] = 1.0
+        out = (input_ > 0).float()
         return out
 
     @staticmethod
@@ -72,8 +71,7 @@ class Sigmoid(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input_):
         ctx.save_for_backward(input_)
-        out = torch.zeros_like(input_)
-        out[input_ > 0] = 1.0
+        out = (input_ > 0).float()
         return out
 
     @staticmethod
@@ -119,8 +117,7 @@ def __init__(self, beta=1):
     @staticmethod
     def forward(ctx, input_):
         ctx.save_for_backward(input_)
-        out = torch.zeros_like(input_)
-        out[input_ > 0] = 1.0
+        out = (input_ > 0).float()
         return out
 
     def backward(self, ctx, grad_output):
diff --git a/tox.ini b/tox.ini
index db37ea03..a6558eaf 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,7 +1,7 @@
 [tox]
 envlist = py36, py37, py38, flake8
 
-[travis]
+[gh-actions]
 python =
     3.8: py38
     3.7: py37


Test patch:
diff --git a/.github/workflows/python-publish-test.yml b/.github/workflows/python-publish-test.yml
new file mode 100644
index 00000000..38e1693e
--- /dev/null
+++ b/.github/workflows/python-publish-test.yml
@@ -0,0 +1,47 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: Upload Python Package to Test PyPi
+
+on:
+  workflow_run: 
+    workflows: ["Python package tag"]
+    # branches: [main, git-actions] 
+    types:
+      - completed
+
+
+jobs:
+
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine
+
+
+    
+    # Upload to TestPyPI
+    - name: Build and Upload to TestPyPI
+      run: |
+        python3 setup.py sdist bdist_wheel
+        python3 -m twine upload dist/*
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.TWINE_TEST_TOKEN }}
+        TWINE_REPOSITORY: testpypi
+
+    # - name: Build and publish
+    #   env:
+    #     TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+    #     TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+    #   run: |
+    #     python setup.py sdist bdist_wheel
+    #     twine upload dist/*
\ No newline at end of file