[logging] Prevent PL from adding dataloader index to logged metrics

src/mednet/engine/callbacks.py

@@ -362,11 +362,17 @@ class LoggingCallback(lightning.pytorch.Callback):
             out which dataset was used for this validation epoch.
         """
 
+        if dataloader_idx == 0:
+            key = "loss/validation"
+        else:
+            key = f"loss/validation-{dataloader_idx}"
+
         pl_module.log(
-            "loss/validation",
+            key,
             outputs.item(),
             prog_bar=False,
             on_step=False,
             on_epoch=True,
             batch_size=batch[0].shape[0],
+            add_dataloader_idx=False,
         )

src/mednet/engine/trainer.py

@@ -94,8 +94,6 @@ def run(
     )
 
     monitor_key = "loss/validation"
-    if len(datamodule.val_dataset_keys()) > 1:
-        monitor_key = "loss/validation/dataloader_idx_0"
 
     # This checkpointer will operate at the end of every validation epoch
     # (which happens at each checkpoint period), it will then save the lowest