🐛 Try to fix callback error

boostcampaitech7 · Sep 15, 2024 · 53a5d12 · 53a5d12
1 parent 8931236
commit 53a5d12
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/engine/callbacks.py b/engine/callbacks.py
@@ -20,3 +20,10 @@ def on_test_end(self, *args, **kwargs):
         file_name = f"{self.model_name}_{date.today()}.csv"
         test_info.to_csv(file_name, index=False, lineterminator='\n')
         print("Output csv file successfully saved!!")
+
+
+    def load_state_dict(self, state_dict):
+        self.predictions = state_dict['predictions']
+
+    def state_dict(self):
+        return {'predictions': self.predictions}