Not using seq_len

[TaylanSoydan]

🐛 Describe the bug

usage.py isnt working on a sequence, seq_len is not used in any functions, only batch dim and input dim are used. there seems to be something fundamental missing in the code.

Versions

Python 3.9.19
absl-py==2.1.0
aiohappyeyeballs==2.4.0
aiohttp==3.10.5
aiosignal==1.3.1
antlr4-python3-runtime==4.9.3
astunparse==1.6.3
async-timeout==4.0.3
attrs==24.2.0
audioread==3.0.1
certifi==2024.8.30
cffi==1.17.1
charset-normalizer==3.3.2
chex==0.1.86
click==8.1.7
contourpy==1.3.0
cycler==0.12.1
datasets==2.21.0
decorator==5.1.1
dill==0.3.8
dm-haiku==0.0.12
docker-pycreds==0.4.0
einops==0.8.0
etils==1.5.2
expelliarmus==1.1.12
filelock==3.15.4
flatbuffers==24.3.25
flax==0.8.5
fonttools==4.53.1
frozenlist==1.4.1
fsspec==2024.6.1
gast==0.6.0
gitdb==4.0.11
GitPython==3.1.43
google-pasta==0.2.0
grpcio==1.66.1
h5py==3.11.0
huggingface-hub==0.24.6
humanize==4.10.0
hydra-core==1.3.2
idna==3.8
importlib_metadata==8.4.0
importlib_resources==6.4.4
importRosbag==1.0.4
jax==0.4.30
jax-cuda12-pjrt==0.4.30
jax-cuda12-plugin==0.4.30
jaxlib==0.4.30
jaxtyping==0.2.34
Jinja2==3.1.4
jmp==0.0.4
joblib==1.4.2
keras==3.5.0
kiwisolver==1.4.7
lazy_loader==0.4
libclang==18.1.1
librosa==0.10.2.post1
llvmlite==0.43.0
Markdown==3.7
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.9.2
mdurl==0.1.2
ml-dtypes==0.4.0
mpmath==1.3.0
msgpack==1.0.8
multidict==6.0.5
multiprocess==0.70.16
namex==0.0.8
nest-asyncio==1.6.0
networkx==3.2.1
numba==0.60.0
numpy==1.26.4
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvcc-cu12==12.6.68
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.20.5
nvidia-nvjitlink-cu12==12.6.68
nvidia-nvtx-cu12==12.1.105
omegaconf==2.3.0
opt-einsum==3.3.0
optax==0.2.3
optree==0.12.1
orbax-checkpoint==0.6.1
packaging==24.1
pandas==2.2.2
pbr==6.1.0
pillow==10.4.0
platformdirs==4.2.2
pooch==1.8.2
promise==2.3
protobuf==3.20.3
psutil==6.0.0
pyarrow==17.0.0
pycparser==2.22
pydub==0.25.1
Pygments==2.18.0
pyparsing==3.1.4
python-dateutil==2.9.0.post0
pytz==2024.1
PyYAML==6.0.2
requests==2.32.3
rich==13.8.0
scikit-base==0.9.0
scikit-learn==1.5.1
scipy==1.13.1
sentry-sdk==2.13.0
setproctitle==1.3.3
six==1.16.0
sktime==0.33.0
smmap==5.0.1
soundfile==0.12.1
soxr==0.5.0.post1
sympy==1.13.2
tabulate==0.9.0
tensorboard==2.17.1
tensorboard-data-server==0.7.2
tensorflow==2.17.0
tensorflow-datasets==4.5.2
tensorflow-io-gcs-filesystem==0.37.1
tensorflow-metadata==1.15.0
tensorstore==0.1.64
termcolor==2.4.0
threadpoolctl==3.5.0
tonic==1.5.0
toolz==0.12.1
torch==1.11.0
torchaudio==0.11.0
torchdiffeq==0.2.4
torchtext==0.12.0
torchvision==0.12.0
tqdm==4.66.5
triton==3.0.0
typeguard==2.13.3
typing_extensions==4.12.2
tzdata==2024.1
urllib3==2.2.2
wandb==0.17.8
Werkzeug==3.0.4
wrapt==1.16.0
xxhash==3.5.0
yarl==1.9.4
zipp==3.20.1