Fixed EWC bug in importances

ContinualAI · Jul 28, 2022 · 555de3d · 555de3d
1 parent 5064e77
commit 555de3d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/avalanche/training/plugins/ewc.py b/avalanche/training/plugins/ewc.py
@@ -89,7 +89,7 @@ def before_backward(self, strategy, **kwargs):
                     # dynamic models may add new units
                     # new units are ignored by the regularization
                     n_units = saved_param.shape[0]
-                    cur_param = saved_param[:n_units]
+                    cur_param = cur_param[:n_units]
                     penalty += (imp * (cur_param - saved_param).pow(2)).sum()
         elif self.mode == "online":
             prev_exp = exp_counter - 1
@@ -101,7 +101,7 @@ def before_backward(self, strategy, **kwargs):
                 # dynamic models may add new units
                 # new units are ignored by the regularization
                 n_units = saved_param.shape[0]
-                cur_param = saved_param[:n_units]
+                cur_param = cur_param[:n_units]
                 penalty += (imp * (cur_param - saved_param).pow(2)).sum()
         else:
             raise ValueError("Wrong EWC mode.")