Unexpected memory requirement jump

[daeh]

I have a model that fits 4 variables via gradient descent on CPU. When I run it locally (MacOS) and on the openmind HPC cluster (Rocky8), it takes about 5GB of RAM. I added a couple more variables and on MacOS the RAM requirement went up about a GB, and on openmind it jumped to 70 GB. Do you have any suggestions for how I might start to track down the platform-specific issue?

[kach]

That's odd! My first instinct is— are the Python/JAX/memo versions exactly the same across both platforms? If you could make sure that the versions are identical across the two platforms, that would help narrow down the scope of the issue!