clean gc after ckpt to avoid oom (#28)

Author: samsja

src/zeroband/checkpoint.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+import gc
 import multiprocessing
 import os
 import time
@@ -162,6 +163,8 @@ def save(self, ckpt_path: str, remote_ckpt_path: str | None) -> None:
 
         self._logger.info(f"Saved checkpoint to {ckpt_path} in {time.perf_counter() - time_start} seconds")
 
+        gc.collect()  # because we are badass engineer
+
         if remote_ckpt_path is not None:
             self._async_save_remote(ckpt_path, remote_ckpt_path)