The memory consumption is almost same

[DaShenZi721]

I have recently been testing the memory usage of several baselines for text tasks on the Long-Range Arena benchmark. However, I have found that regardless of which baseline I use or the batch size, the memory usage is almost the same(around 160MB), as shown in the figure below.

The code can be found:
https://github.com/google-research/long-range-arena/blob/main/lra_benchmarks/text_classification/train.py

I am using the jax-smi, which actually enables an additional thread that runs the jax.profiler.save_device_memory_profile() command every second. I added the following two lines of code in the main function.

from jax_smi import initialise_tracking
initialise_tracking()

My jax environment is as follows

python 3.8.16
jax 0.2.16
jax-lib 0.1.67+cuda111
flax 0.3.0
tensorflow 2.11.0