Fix calculate_kl

kumar-shridhar · Aug 8, 2020 · 16e2462 · 16e2462
1 parent 81e6a35
commit 16e2462
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/metrics.py b/metrics.py
@@ -24,7 +24,7 @@ def acc(outputs, targets):
     return np.mean(outputs.cpu().numpy().argmax(axis=1) == targets.data.cpu().numpy())
 
 
-def calculate_kl(mu_p, sig_p, mu_q, sig_q):
+def calculate_kl(mu_q, sig_q, mu_p, sig_p):
     kl = 0.5 * (2 * torch.log(sig_p / sig_q) - 1 + (sig_q / sig_p).pow(2) + ((mu_p - mu_q) / sig_p).pow(2)).sum()
     return kl