Regularizer funcs are called first with kwargs then positional args

AdvancedPhotonSource · Mar 4, 2024 · 33cf694 · 33cf694
1 parent dfd25ea
commit 33cf694
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/generic_trainer/trainer.py b/generic_trainer/trainer.py
@@ -430,7 +430,10 @@ def compute_losses(self, loss_records, preds, labels):
             pred_dict = self.get_pred_dict(preds)
             for i in range(len(preds), len(self.loss_criterion)):
                 this_loss_func = self.loss_criterion[i]
-                this_loss_tensor = this_loss_func(**pred_dict)
+                try:
+                    this_loss_tensor = this_loss_func(**pred_dict)
+                except TypeError:
+                    this_loss_tensor = this_loss_func(*preds)
                 total_loss_tensor = total_loss_tensor + this_loss_tensor
                 loss_records[i + 1] += this_loss_tensor.detach().item()
         loss_records[0] += total_loss_tensor.detach().item()