s

Author: kddubey

src/grouping_trainer/loss.py

@@ -48,7 +48,8 @@ def _mrl_loss(
         # 1. Detach embeddings from graph
         # 2. Loop over dims, compute loss, backward in loop
         # 3. Backprop detached embeddings' gradients to model
-    return loss_total / len(dim_indices)
+    weight_total = sum(mrl_dim_to_weight[mrl_dims[idx]] for idx in dim_indices)
+    return loss_total / weight_total
 
 
 class PairwiseLoss(torch.nn.Module, ABC):