diff --git a/nir/ir/neuron.py b/nir/ir/neuron.py
index f397475..be94049 100644
--- a/nir/ir/neuron.py
+++ b/nir/ir/neuron.py
@@ -60,8 +60,8 @@ def __post_init__(self):
         ), "All parameters must have the same shape"
         # If w_in is a scalar, make it an array of same shape as v_threshold
         self.w_in = np.ones_like(self.v_threshold) * self.w_in
-        self.input_type = {"input": np.array(self.v_threshold.shape)}
-        self.output_type = {"output": np.array(self.v_threshold.shape)}
+        self.input_type = {"input": np.array([np.size(self.v_threshold)])}
+        self.output_type = {"output": np.array([np.size(self.v_threshold)])}
 
 
 @dataclass(eq=False)
diff --git a/paper/03_rnn/rockpool_apply.ipynb b/paper/03_rnn/rockpool_apply.ipynb
index 966266c..f0f948b 100644
--- a/paper/03_rnn/rockpool_apply.ipynb
+++ b/paper/03_rnn/rockpool_apply.ipynb
@@ -107,21 +107,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 34,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 140/140 [00:13<00:00, 10.58it/s]"
+      "100%|██████████| 140/140 [00:13<00:00, 10.20it/s]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "n_correct 100 out of 140 (71.42857142857143%)\n"
+      "n_correct 108 out of 140 (77.14285714285715%)\n"
      ]
     },
     {
@@ -143,7 +143,7 @@
     "    sample = single_sample[0].numpy()[0]  # shape: (256, 12)\n",
     "\n",
     "    output, rec_dict = net(torch.tensor(sample))\n",
-    "    n_output_spikes = np.sum(output.detach().numpy(), axis=1)\n",
+    "    n_output_spikes = np.sum(output.squeeze().detach().numpy(), axis=0)\n",
     "\n",
     "    predicted_label = int(np.argmax(n_output_spikes))\n",
     "    actual_label = int(single_sample[1])\n",
@@ -160,7 +160,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 35,
    "metadata": {},
    "outputs": [
     {