Out of memory when trainer.save_checkpoint("example.ckpt")
The sbatch session crashes and I get the following error when I include trainer.save_checkpoint("example.ckpt") in my code.

/var/spool/slurmd/job220424/slurm_script: line 15: 39865 Killed                  python roco_train_mlm_lightning.py --run_name debug --precision 16 --mlm_prob 0.15
slurmstepd: error: Detected 3 oom-kill event(s) in step 220424.batch cgroup. Some of your processes may have been killed by the cgroup out-of-memory handler.

Same thing happens when I run the code in notebook on the remoteserver. The kernel dies. Please help. Thank you.