Script for computing memory consumption

[maximzubkov]

Hello again!

I succeeded to run your code on my devices, but now I am struggling with computing memory consumption. Could you please opensource script for computing it?

[MostafaDehghani]

Unfortunately, the memory consumptions reported in the paper are calculated using a google internal tool which is not open-sourced. We will check with the internal teams to see if there is an alternative open-sourced version of that available.