diff --git a/examples/MNIST Addition - Video.ipynb b/examples/MNIST Addition - Video.ipynb
index 62ffe92..4eaad31 100644
--- a/examples/MNIST Addition - Video.ipynb	
+++ b/examples/MNIST Addition - Video.ipynb	
@@ -16,16 +16,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/hanliying/opt/anaconda3/envs/pytorch/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "<torch._C.Generator at 0x7fe93808f970>"
+       "<torch._C.Generator at 0x7fabfcfe9030>"
       ]
      },
-     "execution_count": 2,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -47,7 +55,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -86,9 +94,92 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./MNIST/MNIST/raw/train-images-idx3-ubyte.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "9913344it [00:01, 6405934.93it/s]                              \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extracting ./MNIST/MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/MNIST/raw\n",
+      "\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/MNIST/raw/train-labels-idx1-ubyte.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "29696it [00:00, 7006866.09it/s]          "
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extracting ./MNIST/MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/MNIST/raw\n",
+      "\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "1649664it [00:00, 18619900.63it/s]         \n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extracting ./MNIST/MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/MNIST/raw\n",
+      "\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
+      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "5120it [00:00, 2585772.00it/s]          "
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extracting ./MNIST/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/MNIST/raw\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
    "source": [
     "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n",
     "mnist_train_data = torchvision.datasets.MNIST(root='./MNIST', train=True, download=True,transform=transform)\n",
@@ -107,7 +198,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -145,7 +236,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -162,7 +253,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -189,7 +280,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -199,7 +290,7 @@
     "\n",
     "from pylon.constraint import constraint\n",
     "\n",
-    "def enforce_sum(img1, img2, kwargs):\n",
+    "def enforce_sum(img1, img2, **kwargs):\n",
     "    return img1 + img2 == kwargs['summation']\n",
     "\n",
     "\n",
@@ -215,7 +306,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "metadata": {
     "scrolled": false
    },
@@ -224,14 +315,16 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 11%|█         | 1003/9000 [00:54<21:14,  6.27it/s]"
+      "  0%|          | 0/9000 [00:00<?, ?it/s]/Users/hanliying/Documents/UCLA/Research/pylon/pylon/examples/../pylon/brute_force_solver.py:50: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
+      "  else torch.tensor(data=self.cond(*sample, **kwargs), dtype=torch.bool) for sample in samples ])\n",
+      " 11%|█         | 1005/9000 [00:54<25:23,  5.25it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 6282/10000 (63%)\n",
+      "Test set: Accuracy: 6328/10000 (63%)\n",
       "\n"
      ]
     },
@@ -239,14 +332,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 22%|██▏       | 2004/9000 [01:50<25:17,  4.61it/s]"
+      " 22%|██▏       | 2010/9000 [01:22<11:46,  9.89it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9333/10000 (93%)\n",
+      "Test set: Accuracy: 9310/10000 (93%)\n",
       "\n"
      ]
     },
@@ -254,14 +347,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 33%|███▎      | 3003/9000 [02:46<20:46,  4.81it/s]"
+      " 33%|███▎      | 3009/9000 [02:01<13:27,  7.42it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9581/10000 (96%)\n",
+      "Test set: Accuracy: 9530/10000 (95%)\n",
       "\n"
      ]
     },
@@ -269,14 +362,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 44%|████▍     | 4003/9000 [03:46<15:12,  5.47it/s]"
+      " 45%|████▍     | 4009/9000 [02:27<11:41,  7.12it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9630/10000 (96%)\n",
+      "Test set: Accuracy: 9650/10000 (96%)\n",
       "\n"
      ]
     },
@@ -284,14 +377,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 56%|█████▌    | 5002/9000 [04:46<16:40,  4.00it/s]"
+      " 56%|█████▌    | 5009/9000 [02:57<11:37,  5.72it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9492/10000 (95%)\n",
+      "Test set: Accuracy: 9472/10000 (95%)\n",
       "\n"
      ]
     },
@@ -299,14 +392,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 67%|██████▋   | 6002/9000 [05:47<12:38,  3.95it/s]"
+      " 67%|██████▋   | 6007/9000 [03:30<08:38,  5.77it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9679/10000 (97%)\n",
+      "Test set: Accuracy: 9681/10000 (97%)\n",
       "\n"
      ]
     },
@@ -314,14 +407,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 78%|███████▊  | 7004/9000 [06:49<06:48,  4.88it/s]"
+      " 78%|███████▊  | 7008/9000 [04:00<04:04,  8.16it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9685/10000 (97%)\n",
+      "Test set: Accuracy: 9659/10000 (97%)\n",
       "\n"
      ]
     },
@@ -329,14 +422,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      " 89%|████████▉ | 8004/9000 [07:51<03:23,  4.88it/s]"
+      " 89%|████████▉ | 8006/9000 [04:26<01:35, 10.37it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9716/10000 (97%)\n",
+      "Test set: Accuracy: 9593/10000 (96%)\n",
       "\n"
      ]
     },
@@ -344,14 +437,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 9000/9000 [08:52<00:00, 16.91it/s]\n"
+      "100%|██████████| 9000/9000 [04:52<00:00, 30.78it/s]\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test set: Accuracy: 9685/10000 (97%)\n",
+      "Test set: Accuracy: 9669/10000 (97%)\n",
       "\n"
      ]
     }
@@ -392,7 +485,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3.8.12 ('pytorch')",
    "language": "python",
    "name": "python3"
   },
@@ -406,7 +499,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.7"
+   "version": "3.8.12"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "7a719078278c4492d805b55cadf911a94c7633b48007fee7a7c47218b923b9ef"
+   }
   }
  },
  "nbformat": 4,
diff --git a/pylon/brute_force_solver.py b/pylon/brute_force_solver.py
index e601c07..1adef5e 100644
--- a/pylon/brute_force_solver.py
+++ b/pylon/brute_force_solver.py
@@ -47,7 +47,7 @@ def loss(self, *logits, **kwargs):
         # Get all possible decodings
         samples = self.all_samples(log_probs)
         indices = torch.stack([torch.tensor(data=self.cond(*sample), dtype=torch.bool) if kwargs == {}
-            else torch.tensor(data=self.cond(*sample), dtype=torch.bool) for sample in samples ])
+            else torch.tensor(data=self.cond(*sample, **kwargs), dtype=torch.bool) for sample in samples ])
 
         losses = torch.stack([decoding_loss(sample, log_probs) for sample in samples])