From 93859ad866675a04e5bc62d5cc953f17958d738f Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Wed, 13 Nov 2024 08:47:00 -0800
Subject: [PATCH] document dimensions

---
 pi_zero_pytorch/pi_zero.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pi_zero_pytorch/pi_zero.py b/pi_zero_pytorch/pi_zero.py
index 16ee241..a0b32ed 100644
--- a/pi_zero_pytorch/pi_zero.py
+++ b/pi_zero_pytorch/pi_zero.py
@@ -29,6 +29,7 @@
 
 # b - batch
 # n - sequence
+# i, j - attention row and col sequence
 # na - seq of actions
 # nt - seq of text tokens
 # nv - seq of visual tokens