From 3aa2fa96eef0e11d7aabbe13d838a67a5c003c45 Mon Sep 17 00:00:00 2001
From: Mark Kurtz <mark@neuralmagic.com>
Date: Wed, 12 May 2021 17:11:00 -0400
Subject: [PATCH 1/3] Fix bugs with ScheduledOptimizer and pytorch scheduler
 refactor

---
 src/sparseml/pytorch/optim/modifier.py  |  4 +---
 src/sparseml/pytorch/optim/optimizer.py | 22 +++++++++++-----------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/src/sparseml/pytorch/optim/modifier.py b/src/sparseml/pytorch/optim/modifier.py
index 907867c583d..55ac01bace4 100644
--- a/src/sparseml/pytorch/optim/modifier.py
+++ b/src/sparseml/pytorch/optim/modifier.py
@@ -196,9 +196,7 @@ def initialize(
             for individual modifiers.
         """
         self._initialized = True
-
-        if loggers:
-            self.initialize_loggers(loggers)
+        self.initialize_loggers(loggers)
 
     def initialize_loggers(self, loggers: Union[None, List[BaseLogger]]):
         """
diff --git a/src/sparseml/pytorch/optim/optimizer.py b/src/sparseml/pytorch/optim/optimizer.py
index 02a2def1eb5..cdb9788e7de 100644
--- a/src/sparseml/pytorch/optim/optimizer.py
+++ b/src/sparseml/pytorch/optim/optimizer.py
@@ -81,11 +81,11 @@ def __init__(
     ):
         # do not call into super since this instance is not passing all calls to
         # the nested optimizer
-        warnings.warn(
-            "ScheduledOptimizer is deprecated and will be deleted in the future. "
-            "Please replace with manager.modify",
-            UserWarning,
-        )
+        # warnings.warn(
+        #     "ScheduledOptimizer is deprecated and will be deleted in the future. "
+        #     "Please replace with manager.modify",
+        #     UserWarning,
+        # )  TODO: uncomment in next release once docs are ready
 
         manager.initialize(module, epoch=0.0, loggers=loggers)
         self._wrapper = RecipeManagerStepWrapper(
@@ -107,7 +107,7 @@ def __getattr__(self, item):
         if item in self.__dict__:
             return getattr(self, item)
 
-        return getattr(self._wrapped, item)
+        return getattr(self._wrapper.wrapped_optimizer, item)
 
     def __setattr__(self, key, value):
         if key in [
@@ -118,7 +118,7 @@ def __setattr__(self, key, value):
         ]:
             super().__setattr__(key, value)
         else:
-            setattr(self._optimizer, key, value)
+            setattr(self._wrapper.wrapped_optimizer, key, value)
 
     @property
     def learning_rate(self) -> float:
@@ -126,7 +126,7 @@ def learning_rate(self) -> float:
         :return: convenience function to get the first learning rate for any of
             the param groups in the optimizer
         """
-        return get_optim_learning_rate(self._optimizer)
+        return get_optim_learning_rate(self._wrapper.wrapped_optimizer)
 
     @learning_rate.setter
     def learning_rate(self, value: float):
@@ -134,7 +134,7 @@ def learning_rate(self, value: float):
         :param value: the learning rate to set for the optimizer,
             will set all param groups in the optim to this value
         """
-        set_optim_learning_rate(self._optimizer, value)
+        set_optim_learning_rate(self._wrapper.wrapped_optimizer, value)
 
     @property
     def manager(self) -> ScheduledModifierManager:
@@ -144,10 +144,10 @@ def manager(self) -> ScheduledModifierManager:
         return self._wrapper.wrapped_manager
 
     def manager_state_dict(self):
-        return self._manager.state_dict()
+        return self._wrapper.wrapped_manager.state_dict()
 
     def load_manager_state_dict(self, state_dict):
-        self._manager.load_state_dict(state_dict)
+        self._wrapper.wrapped_manager.load_state_dict(state_dict)
 
     def step(self, closure=None):
         """

From d5c7314de1f6feb4d83ae13cec17d2c01207d7ed Mon Sep 17 00:00:00 2001
From: Mark Kurtz <mark@neuralmagic.com>
Date: Wed, 12 May 2021 17:18:50 -0400
Subject: [PATCH 2/3] fix test cases for updated log_update initialization

---
 notebooks/pytorch_classification.ipynb        | 355 +++++++++++++++++-
 tests/sparseml/pytorch/optim/test_modifier.py |  10 -
 2 files changed, 345 insertions(+), 20 deletions(-)

diff --git a/notebooks/pytorch_classification.ipynb b/notebooks/pytorch_classification.ipynb
index cdb71ebb0c9..255bc5c638e 100644
--- a/notebooks/pytorch_classification.ipynb
+++ b/notebooks/pytorch_classification.ipynb
@@ -34,7 +34,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -62,9 +62,284 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "loading model...\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b79a138cd5f24ec685ebe1f2a2a4eaa1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "downloading...:   0%|          | 0/94371997 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ResNet(\n",
+      "  (input): _Input(\n",
+      "    (conv): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
+      "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "    (act): ReLU(inplace=True)\n",
+      "    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
+      "  )\n",
+      "  (sections): Sequential(\n",
+      "    (0): Sequential(\n",
+      "      (0): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (identity): _IdentityModifier(\n",
+      "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        )\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (1): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (2): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (1): Sequential(\n",
+      "      (0): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (identity): _IdentityModifier(\n",
+      "          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+      "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        )\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (1): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (2): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (3): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (2): Sequential(\n",
+      "      (0): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (identity): _IdentityModifier(\n",
+      "          (conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+      "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        )\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (1): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (2): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (3): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (4): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (5): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (3): Sequential(\n",
+      "      (0): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (identity): _IdentityModifier(\n",
+      "          (conv): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
+      "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        )\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (1): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "      (2): _BottleneckBlock(\n",
+      "        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act1): ReLU(inplace=True)\n",
+      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
+      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act2): ReLU(inplace=True)\n",
+      "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "        (act_out): ReLU(inplace=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (classifier): _Classifier(\n",
+      "    (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
+      "    (fc): Linear(in_features=2048, out_features=10, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      ")\n",
+      "\n",
+      "loading train dataset...\n",
+      "already downloaded imagenette ImagenetteSize.s320\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Dataset ImagenetteDataset\n",
+      "    Number of datapoints: 12894\n",
+      "    Root location: /Users/markkurtz/.cache/nm_datasets/imagenette/imagenette-320/train\n",
+      "    StandardTransform\n",
+      "Transform: Compose(\n",
+      "               Resize(size=256, interpolation=bilinear)\n",
+      "               CenterCrop(size=(224, 224))\n",
+      "               ToTensor()\n",
+      "               Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
+      "           )\n",
+      "\n",
+      "loading val dataset...\n",
+      "already downloaded imagenette ImagenetteSize.s320\n",
+      "Dataset ImagenetteDataset\n",
+      "    Number of datapoints: 500\n",
+      "    Root location: /Users/markkurtz/.cache/nm_datasets/imagenette/imagenette-320/val\n",
+      "    StandardTransform\n",
+      "Transform: Compose(\n",
+      "               Resize(size=256, interpolation=bilinear)\n",
+      "               CenterCrop(size=(224, 224))\n",
+      "               ToTensor()\n",
+      "               Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
+      "           )\n"
+     ]
+    }
+   ],
    "source": [
     "from sparseml.pytorch.models import ModelRegistry\n",
     "from sparseml.pytorch.datasets import ImagenetteDataset, ImagenetteSize\n",
@@ -109,7 +384,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -163,9 +438,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cpu\n"
+     ]
+    }
+   ],
    "source": [
     "from torch.utils.data import DataLoader\n",
     "from torch.nn import CrossEntropyLoss\n",
@@ -215,9 +498,31 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2e001ac45145492ca703fdd125d37953",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "downloading...:   0%|          | 0/4210 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Recipe downloaded to: /Users/markkurtz/.cache/sparsezoo/d863d8ed-be54-4b9e-8a08-6bdec8219ec5/recipes/original.md\n"
+     ]
+    }
+   ],
    "source": [
     "from sparsezoo import Zoo\n",
     "\n",
@@ -240,7 +545,37 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running Training Epoch 1/10\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/markkurtz/code/github/neuralmagic/sparseml/src/sparseml/pytorch/optim/optimizer.py:87: UserWarning: ScheduledOptimizer is deprecated and will be deleted in the future. Please replace with manager.modify\n",
+      "  UserWarning,\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f91357fba3ea484889e6c6c8dbb34301",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/101 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "from sparseml.pytorch.optim import (\n",
     "    ScheduledModifierManager,\n",
@@ -359,7 +694,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.9"
+   "version": "3.6.8"
   }
  },
  "nbformat": 4,
diff --git a/tests/sparseml/pytorch/optim/test_modifier.py b/tests/sparseml/pytorch/optim/test_modifier.py
index c0b05ca10fc..2ceb55c7dd0 100644
--- a/tests/sparseml/pytorch/optim/test_modifier.py
+++ b/tests/sparseml/pytorch/optim/test_modifier.py
@@ -216,11 +216,6 @@ def test_log_update(
         model = model_lambda()
         optimizer = optim_lambda(model)
 
-        with pytest.raises(RuntimeError):
-            modifier.log_update(model, optimizer, test_epoch, test_steps_per_epoch)
-
-        self.initialize_helper(modifier, model, log_initialize=False)
-
         with pytest.raises(RuntimeError):
             modifier.log_update(model, optimizer, test_epoch, test_steps_per_epoch)
 
@@ -496,11 +491,6 @@ def test_scheduled_log_update(
         model = model_lambda()
         optimizer = optim_lambda(model)
 
-        with pytest.raises(RuntimeError):
-            modifier.scheduled_log_update(model, optimizer, 0.0, test_steps_per_epoch)
-
-        self.initialize_helper(modifier, model, log_initialize=False)
-
         with pytest.raises(RuntimeError):
             modifier.scheduled_log_update(model, optimizer, 0.0, test_steps_per_epoch)
 

From 4aca5a850e15c0ed5a0eb2eb45788ce5f7cdb49d Mon Sep 17 00:00:00 2001
From: Mark Kurtz <mark@neuralmagic.com>
Date: Wed, 12 May 2021 17:20:53 -0400
Subject: [PATCH 3/3] revert changes to classification notebook

---
 notebooks/pytorch_classification.ipynb | 355 +------------------------
 1 file changed, 10 insertions(+), 345 deletions(-)

diff --git a/notebooks/pytorch_classification.ipynb b/notebooks/pytorch_classification.ipynb
index 255bc5c638e..cdb71ebb0c9 100644
--- a/notebooks/pytorch_classification.ipynb
+++ b/notebooks/pytorch_classification.ipynb
@@ -34,7 +34,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -62,284 +62,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "loading model...\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b79a138cd5f24ec685ebe1f2a2a4eaa1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "downloading...:   0%|          | 0/94371997 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "ResNet(\n",
-      "  (input): _Input(\n",
-      "    (conv): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
-      "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "    (act): ReLU(inplace=True)\n",
-      "    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
-      "  )\n",
-      "  (sections): Sequential(\n",
-      "    (0): Sequential(\n",
-      "      (0): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (identity): _IdentityModifier(\n",
-      "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        )\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (1): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (2): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "    )\n",
-      "    (1): Sequential(\n",
-      "      (0): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (identity): _IdentityModifier(\n",
-      "          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
-      "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        )\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (1): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (2): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (3): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "    )\n",
-      "    (2): Sequential(\n",
-      "      (0): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (identity): _IdentityModifier(\n",
-      "          (conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
-      "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        )\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (1): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (2): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (3): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (4): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (5): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "    )\n",
-      "    (3): Sequential(\n",
-      "      (0): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (identity): _IdentityModifier(\n",
-      "          (conv): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
-      "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        )\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (1): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "      (2): _BottleneckBlock(\n",
-      "        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act1): ReLU(inplace=True)\n",
-      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
-      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act2): ReLU(inplace=True)\n",
-      "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
-      "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-      "        (act_out): ReLU(inplace=True)\n",
-      "      )\n",
-      "    )\n",
-      "  )\n",
-      "  (classifier): _Classifier(\n",
-      "    (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
-      "    (fc): Linear(in_features=2048, out_features=10, bias=True)\n",
-      "    (softmax): Softmax(dim=1)\n",
-      "  )\n",
-      ")\n",
-      "\n",
-      "loading train dataset...\n",
-      "already downloaded imagenette ImagenetteSize.s320\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Dataset ImagenetteDataset\n",
-      "    Number of datapoints: 12894\n",
-      "    Root location: /Users/markkurtz/.cache/nm_datasets/imagenette/imagenette-320/train\n",
-      "    StandardTransform\n",
-      "Transform: Compose(\n",
-      "               Resize(size=256, interpolation=bilinear)\n",
-      "               CenterCrop(size=(224, 224))\n",
-      "               ToTensor()\n",
-      "               Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
-      "           )\n",
-      "\n",
-      "loading val dataset...\n",
-      "already downloaded imagenette ImagenetteSize.s320\n",
-      "Dataset ImagenetteDataset\n",
-      "    Number of datapoints: 500\n",
-      "    Root location: /Users/markkurtz/.cache/nm_datasets/imagenette/imagenette-320/val\n",
-      "    StandardTransform\n",
-      "Transform: Compose(\n",
-      "               Resize(size=256, interpolation=bilinear)\n",
-      "               CenterCrop(size=(224, 224))\n",
-      "               ToTensor()\n",
-      "               Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
-      "           )\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from sparseml.pytorch.models import ModelRegistry\n",
     "from sparseml.pytorch.datasets import ImagenetteDataset, ImagenetteSize\n",
@@ -384,7 +109,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -438,17 +163,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Using device: cpu\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from torch.utils.data import DataLoader\n",
     "from torch.nn import CrossEntropyLoss\n",
@@ -498,31 +215,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2e001ac45145492ca703fdd125d37953",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "downloading...:   0%|          | 0/4210 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Recipe downloaded to: /Users/markkurtz/.cache/sparsezoo/d863d8ed-be54-4b9e-8a08-6bdec8219ec5/recipes/original.md\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from sparsezoo import Zoo\n",
     "\n",
@@ -545,37 +240,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Running Training Epoch 1/10\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/markkurtz/code/github/neuralmagic/sparseml/src/sparseml/pytorch/optim/optimizer.py:87: UserWarning: ScheduledOptimizer is deprecated and will be deleted in the future. Please replace with manager.modify\n",
-      "  UserWarning,\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f91357fba3ea484889e6c6c8dbb34301",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/101 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "from sparseml.pytorch.optim import (\n",
     "    ScheduledModifierManager,\n",
@@ -694,7 +359,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.8"
+   "version": "3.6.9"
   }
  },
  "nbformat": 4,