unit test passed

marrlab · Oct 14, 2024 · 141fb6f · 141fb6f
1 parent 80bc363
commit 141fb6f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/domainlab/algos/trainers/train_miro.py b/domainlab/algos/trainers/train_miro.py
@@ -58,5 +58,5 @@ def _cal_reg_loss(self, tensor_x, tensor_y, tensor_d, others=None):
             mean_ref = list_batch_inter_feat_ref[ind_layer]
             mean_ref = mean_ref.to(device)
             vlb = (mean - mean_ref).pow(2).div(var) + var.log()
-            reg_loss += vlb / 2.
+            reg_loss += vlb.mean(axis=-1) / 2.
         return [reg_loss], [self.aconf.gamma_reg]