minor

pomonam · Jul 5, 2024 · c78c092 · c78c092
1 parent c56e285
commit c78c092
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/kronfluence/computer/factor_computer.py b/kronfluence/computer/factor_computer.py
@@ -298,6 +298,9 @@ def fit_covariance_matrices(
                         total_data_examples=max_partition_examples,
                     )
 
+                self._reset_memory()
+
+
                 import gc
                 for obj in gc.get_objects():
                     try:
@@ -308,6 +311,7 @@ def fit_covariance_matrices(
                     except:
                         pass
 
+
                 start_time = get_time(state=self.state)
                 with self.profiler.profile("Fit Covariance"):
                     loader = self._get_dataloader(
@@ -332,7 +336,6 @@ def fit_covariance_matrices(
                     f"Fitted covariance matrices with {num_data_processed.item()} data points in "
                     f"{elapsed_time:.2f} seconds."
                 )
-                self._reset_memory()
 
                 print("Done")
                 import gc