Repository: continualai/avalanche
Issue: The LwFPlugin return a negative loss

🐛 **Describe the bug**
I noticed the loss in tensorboard log is negative. And only ReplayPlugin and LwFPlugin were used. 

I debug the LwFPlugin code and find that the kl_div return a negative loss.

After the code "[ : . au ]" the sum of log_p or q is not 1, and may be this is a reason. 

🐝 **Expected behavior**
Does the negative loss is corret and make any sense? Or it should be fixed?


🦋 **Additional context**
The LwFPlugin version i used is shown below.

     def _distillation_loss(self, out, prev_out, active_units):
        """
        Compute distillation loss between output of the current model and
        and output of the previous (saved) model.
        """
        # we compute the loss only on the previously active units.
        au = list(active_units)
        log_p = torch.log_softmax(out / self.temperature, dim=1)[:, au]
        q = torch.softmax(prev_out / self.temperature, dim=1)[:, au]
        res = torch.nn.functional.kl_div(log_p, q, reduction="batchmean")        
        res=self.disstill_loss(log_p,q)
        return res

    def penalty(self, out, x, alpha, curr_model):
        """
        Compute weighted distillation loss.
        """
        if self.prev_model is None:
            return 0
        else:
            with torch.no_grad():
                if isinstance(self.prev_model, MultiTaskModule):
                    # output from previous output heads.
                    y_prev = avalanche_forward(self.prev_model, x, None)
                    # in a multitask scenario we need to compute the output
                    # from all the heads, so we need to call forward again.
                    # TODO: can we avoid this?
                    y_curr = avalanche_forward(curr_model, x, None)
                else:  # no task labels
                    y_prev = {"0": self.prev_model(x)}
                    y_curr = {"0": out}

            dist_loss = 0
            for task_id in y_prev.keys():
                # compute kd only for previous heads.
                if task_id in self.prev_classes:
                    yp = y_prev[task_id]
                    yc = y_curr[task_id]
                    au = self.prev_classes[task_id]
                    dist_loss += self._distillation_loss(yc, yp, au)
            return alpha * dist_loss

Changed files:
avalanche/evaluation/metrics/mean_scores.py, avalanche/evaluation/metrics/labels_repartition.py, avalanche/training/plugins/lwf.py, avalanche/training/supervised/ar1.py

diff --git a/avalanche/evaluation/metrics/labels_repartition.py b/avalanche/evaluation/metrics/labels_repartition.py
index bd0cb987..2b85f855 100644
--- a/avalanche/evaluation/metrics/labels_repartition.py
+++ b/avalanche/evaluation/metrics/labels_repartition.py
@@ -40,7 +40,7 @@ class LabelsRepartition(Metric):
         self.class_order = None
         self.reset()
 
-    def reset(self) -> None:
+    def reset(self, **kargs) -> None:
         self.task2label2count = defaultdict(Counter)
 
     def update(
@@ -112,9 +112,9 @@ class LabelsRepartitionPlugin(GenericPluginMetric[Figure]):
             dict
         )
 
-    def reset(self) -> None:
-        self.steps.append(self.global_it_counter)
-        return super().reset()
+    def reset(self, strategy) -> None:
+        self.steps.append(strategy.clock.train_iterations)
+        return super().reset(strategy)
 
     def update(self, strategy: "SupervisedTemplate"):
         if strategy.clock.train_exp_epochs and self.emit_reset_at != "epoch":
@@ -128,7 +128,7 @@ class LabelsRepartitionPlugin(GenericPluginMetric[Figure]):
         )
 
     def _package_result(self, strategy: "SupervisedTemplate") -> "MetricResult":
-        self.steps.append(self.global_it_counter)
+        self.steps.append(strategy.clock.train_iterations)
         task2label2count = self.labels_repartition.result()
         for task, label2count in task2label2count.items():
             for label, count in label2count.items():
diff --git a/avalanche/evaluation/metrics/mean_scores.py b/avalanche/evaluation/metrics/mean_scores.py
index 2bda93c8..bdc422f8 100644
--- a/avalanche/evaluation/metrics/mean_scores.py
+++ b/avalanche/evaluation/metrics/mean_scores.py
@@ -186,9 +186,10 @@ class MeanScoresPluginMetricABC(PluginMetric, ABC):
 
     def _package_result(self, strategy: "SupervisedTemplate") -> "MetricResult":
         label_cat2mean_score: Dict[LabelCat, float] = self.result()
+        num_it = strategy.clock.train_iterations
 
         for label_cat, m in label_cat2mean_score.items():
-            self.label_cat2step2mean[label_cat][self.global_it_counter] = m
+            self.label_cat2step2mean[label_cat][num_it] = m
 
         base_metric_name = get_metric_name(
             self, strategy, add_experience=False, add_task=False
@@ -199,7 +200,7 @@ class MeanScoresPluginMetricABC(PluginMetric, ABC):
                 self,
                 name=base_metric_name + f"/{label_cat}_classes",
                 value=m,
-                x_plot=self.global_it_counter,
+                x_plot=num_it,
             )
             for label_cat, m in label_cat2mean_score.items()
         ]
@@ -210,7 +211,7 @@ class MeanScoresPluginMetricABC(PluginMetric, ABC):
                     name=base_metric_name + f"/new_old_diff",
                     value=label_cat2mean_score["new"]
                     - label_cat2mean_score["old"],
-                    x_plot=self.global_it_counter,
+                    x_plot=num_it,
                 )
             )
         if self.image_creator is not None:
@@ -222,7 +223,7 @@ class MeanScoresPluginMetricABC(PluginMetric, ABC):
                         self.image_creator(self.label_cat2step2mean),
                         self.label_cat2step2mean,
                     ),
-                    x_plot=self.global_it_counter,
+                    x_plot=num_it,
                 )
             )
 
diff --git a/avalanche/training/plugins/lwf.py b/avalanche/training/plugins/lwf.py
index 93814898..13aa7d0e 100644
--- a/avalanche/training/plugins/lwf.py
+++ b/avalanche/training/plugins/lwf.py
@@ -42,8 +42,9 @@ class LwFPlugin(SupervisedPlugin):
         """
         # we compute the loss only on the previously active units.
         au = list(active_units)
-        log_p = torch.log_softmax(out / self.temperature, dim=1)[:, au]
-        q = torch.softmax(prev_out / self.temperature, dim=1)[:, au]
+
+        log_p = torch.log_softmax(out[:, au] / self.temperature, dim=1)
+        q = torch.softmax(prev_out[:, au] / self.temperature, dim=1)
         res = torch.nn.functional.kl_div(log_p, q, reduction="batchmean")
         return res
 
diff --git a/avalanche/training/supervised/ar1.py b/avalanche/training/supervised/ar1.py
index 0e3de3ef..3a1955c8 100644
--- a/avalanche/training/supervised/ar1.py
+++ b/avalanche/training/supervised/ar1.py
@@ -266,6 +266,7 @@ class AR1(SupervisedTemplate):
 
     def training_epoch(self, **kwargs):
         for mb_it, self.mbatch in enumerate(self.dataloader):
+            self._unpack_minibatch()
             self._before_training_iteration(**kwargs)
 
             self.optimizer.zero_grad()