Merge branch 'main' into pablo-garay/remove_exportDeploy

Author: pablo-garay

nemo/collections/asr/models/sortformer_diar_models.py

@@ -112,8 +112,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
             self.sortformer_modules.encoder_proj = None
         self._init_loss_weights()
 
-        self.eps = 1e-3
-        self.negative_init_val = -99
+        self.eps = self._cfg.get("eps", 1e-3)
+        self.negative_init_val = self._cfg.get("negative_init_val", -99)
         self.loss = instantiate(self._cfg.loss)
 
         self.async_streaming = self._cfg.get("async_streaming", False)
@@ -832,6 +832,7 @@ def _get_aux_train_evaluations(self, preds, targets, target_lens) -> dict:
         Returns:
             (dict): A dictionary containing the following training metrics.
         """
+        targets = targets.to(preds.dtype)
         if preds.shape[1] < targets.shape[1]:
             logging.info(
                 f"WARNING! preds has less frames than targets ({preds.shape[1]} < {targets.shape[1]}). "
@@ -904,6 +905,7 @@ def _get_aux_validation_evaluations(self, preds, targets, target_lens) -> dict:
         Returns:
             val_metrics (dict): A dictionary containing the following validation metrics
         """
+        targets = targets.to(preds.dtype)
         if preds.shape[1] < targets.shape[1]:
             logging.info(
                 f"WARNING! preds has less frames than targets ({preds.shape[1]} < {targets.shape[1]}). "
@@ -1035,6 +1037,7 @@ def _get_aux_test_batch_evaluations(self, batch_idx: int, preds, targets, target
             target_lens (torch.Tensor): Lengths of target sequences.
                 Shape: (batch_size,)
         """
+        targets = targets.to(preds.dtype)
         if preds.shape[1] < targets.shape[1]:
             logging.info(
                 f"WARNING! preds has less frames than targets ({preds.shape[1]} < {targets.shape[1]}). "