more verbose

examples/02_inference.ipynb

@@ -355,9 +355,9 @@
     "- `predictions` (torch.Tensor): logits at each internal tick (Shape: B, C, T)\n",
     "- `certainties` (torch.Tensor): 1 - normalized entropy at each internal tick (Shape: B, 2, T)\n",
     "- `synchronization` (tuple(np.array)): a tuple containing the synch represenations over time for both actions and outputs (Shape of each: T, B, synch_rep_size)\n",
-    "- `pre_activations` (np.array): the history of neuron pre-activations (after synapses, before NLMs)\n",
-    "- `post_activations` (np.array): the history of neuron post-activations (after NLMs)\n",
-    "- `attention_tracking` (np.array): a tensor containing the attention weights\n"
+    "- `pre_activations` (np.array): the history of neuron pre-activations (after synapses, before NLMs) (Shape: T, B, d_model)\n",
+    "- `post_activations` (np.array): the history of neuron post-activations (after NLMs) (Shape: T, B, d_model)\n",
+    "- `attention_tracking` (np.array): a tensor containing the attention weights (Shape: T, B, num_heads, target_seq_length=1, source_seq_length=196)\n"
    ]
   },
   {
@@ -409,6 +409,14 @@
     "print(f\"Predicted Class (most certain): {prediction_most_certain.item()} = {IMAGENET_CLASS_LIST[prediction.item()]}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "fef2a65e",
+   "metadata": {},
+   "source": [
+    "Next, we write a function to make pretty gifs."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,