Running out of memory in the summary example

I am trying to run this notebook locally: notebooks/examples/summarization.ipynb at main · huggingface/notebooks · GitHub
I am running on mac m2 with python 3.14.
The only change I did is to install additional dependencies (pip install datasets transformers torch torchvision torchaudio accelerate).
During the train process I see the memory goes up like a memory leak. After some run time (something like 1300 steps) I get an error of out of memory for the device.
“RuntimeError: MPS backend out of memory (MPS allocated: 4.20 GiB, other allocations: 43.49 GiB, max allowed: 47.74 GiB). Tried to allocate 51.44 MiB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).”
I tried to run torch.mps.empty_cache() but it didn’t change
What can be my issue? How can I fix this memory issue so I will be able to train the model?
I don’t want to set PYTORCH_MPS_HIGH_WATERMARK_RATIO to 0 as it seems like just a hack, and to fix my core issue about the memory leak
Thank for you helping!

(post deleted by author)