Fix error in old PyTorch for KL and JS divergence

ignite/metrics/js_divergence.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn.functional as F
+from packaging.version import Version
 
 from ignite.exceptions import NotComputableError
 from ignite.metrics.kl_divergence import KLDivergence
@@ -71,14 +72,24 @@ class JSDivergence(KLDivergence):
     """
 
     def _update(self, y_pred: torch.Tensor, y: torch.Tensor) -> None:
-        m_prob = (F.softmax(y_pred, dim=1) + F.softmax(y, dim=1)) / 2
+        y_pred_prob = F.softmax(y_pred, dim=1)
+        y_prob = F.softmax(y, dim=1)
+        m_prob = (y_pred_prob + y_prob) / 2
         m_log = m_prob.log()
-        y_pred = F.log_softmax(y_pred, dim=1)
-        y = F.log_softmax(y, dim=1)
-        self._sum_of_kl += (
-            F.kl_div(m_log, y_pred, log_target=True, reduction="sum")
-            + F.kl_div(m_log, y, log_target=True, reduction="sum")
-        ).to(self._device)
+
+        if Version(torch.__version__) >= Version("1.6.0"):
+            # log_target option can be used from 1.6.0
+            y_pred_log = F.log_softmax(y_pred, dim=1)
+            y_log = F.log_softmax(y, dim=1)
+            self._sum_of_kl += (
+                F.kl_div(m_log, y_pred_log, log_target=True, reduction="sum")
+                + F.kl_div(m_log, y_log, log_target=True, reduction="sum")
+            ).to(self._device)
+        else:
+            # y_pred and y are expected to be probabilities
+            self._sum_of_kl += (
+                F.kl_div(m_log, y_pred_prob, reduction="sum") + F.kl_div(m_log, y_prob, reduction="sum")
+            ).to(self._device)
 
     @sync_all_reduce("_sum_of_kl", "_num_examples")
     def compute(self) -> float:

ignite/metrics/kl_divergence.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn.functional as F
+from packaging.version import Version
 
 from ignite.exceptions import NotComputableError
 from ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce
@@ -91,8 +92,16 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
 
     def _update(self, y_pred: torch.Tensor, y: torch.Tensor) -> None:
         y_pred = F.log_softmax(y_pred, dim=1)
-        y = F.log_softmax(y, dim=1)
-        kl_sum = F.kl_div(y_pred, y, log_target=True, reduction="sum")
+
+        if Version(torch.__version__) >= Version("1.6.0"):
+            # log_target option can be used from 1.6.0
+            y = F.log_softmax(y, dim=1)
+            kl_sum = F.kl_div(y_pred, y, log_target=True, reduction="sum")
+        else:
+            # y is expected to be a probability tensor
+            y = F.softmax(y, dim=1)
+            kl_sum = F.kl_div(y_pred, y, reduction="sum")
+
         self._sum_of_kl += kl_sum.to(self._device)
 
     @sync_all_reduce("_sum_of_kl", "_num_examples")