HEAD requires a lot of memory causing OOM

Hi,

I have problems with OOM using ray and the dashboard shows that Head requires large amounts of memory. What does need Head so much memory for? object_store_memory seems to be unrelated. Instead of sharing my complex code, I share a minimum working example from ray (MNIST), that does not run out of memory, but still requires a lot of memory.

I am using ray 2.8.1 and the MNIST example from: Train a PyTorch Model on Fashion MNIST — Ray 2.8.1

How can I reduce the memory requirements by Head? I get the problems on different systems (Linux, Mac) with and without GPUs.

Thank you very much in advance.