@Neel-Gupta I have access to Pro too. And you are right – lscpu says I have 40 cores available. Strange that cat /proc/cpuinfo says something different
Regarding the other issue that I mentioned before --IO problems-- it seems more like a current issue in datasets. Without these two problems, the memory limit is still problematic, even in the 32 GB setting that Colab Pro provides. In some cases, reloading a checkpoint using transformers consumes a lot of RAM
As you just mentioned before, applying to TRC (TPU Research Cloud) seems to be the best option