Merge pull request #109 from APerezFadon:main

PiperOrigin-RevId: 914461539
Change-Id: I8b1f5e1dfa2d707f1a545c03b8d754ba5cd8fc0f

Author: jsspencer

ferminet/train.py

@@ -1069,6 +1069,7 @@ def learning_rate_schedule(t_: jnp.ndarray) -> jnp.ndarray:
 
       # Checkpointing
       if time.time() - time_of_last_ckpt > cfg.log.save_frequency * 60:
+        writer.flush()
         checkpoint.save(ckpt_save_path, t, data, params, opt_state, mcmc_width)
         time_of_last_ckpt = time.time()
 

ferminet/utils/writers.py

@@ -75,5 +75,9 @@ def write(self, t: int, **data):
     if self._log:
       logging.info('Iteration %s: %s', t, data)
 
+  def flush(self):
+    """Flush the cache whenever a checkpoint is saved."""
+    self._file.flush()
+
   def __exit__(self, exc_type, exc_val, exc_tb):
     self._file.close()