diff --git a/enhancer/loss.py b/enhancer/loss.py index 32b30cf..ec753d4 100644 --- a/enhancer/loss.py +++ b/enhancer/loss.py @@ -3,6 +3,7 @@ import logging import numpy as np import torch import torch.nn as nn +from torchmetrics import ScaleInvariantSignalNoiseRatio from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility @@ -188,4 +189,5 @@ LOSS_MAP = { "si-sdr": Si_SDR, "pesq": Pesq, "stoi": Stoi, + "si-snr": ScaleInvariantSignalNoiseRatio, }