From 26dda619cc2f248ac186fac547b1c4917c2536e8 Mon Sep 17 00:00:00 2001
From: BowenD-UCB <84425382+BowenD-UCB@users.noreply.github.com>
Date: Wed, 7 Feb 2024 15:52:42 -0800
Subject: [PATCH] fixed bug that caused test MAE not saved

---
 chgnet/trainer/trainer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/chgnet/trainer/trainer.py b/chgnet/trainer/trainer.py
index 3f90ac8..82e3b6a 100644
--- a/chgnet/trainer/trainer.py
+++ b/chgnet/trainer/trainer.py
@@ -275,6 +275,7 @@ def train(
                     test_loader, is_test=True, test_result_save_path=None
                 )
             self.training_history[key]["test"] = [test_mae[key] for key in self.targets]
+            self.save(filename=os.path.join(save_dir, file))
 
     def _train(self, train_loader: DataLoader, current_epoch: int) -> dict:
         """Train all data for one epoch.