From 4dbefd51b38ca3e4678951d412ce766515b00bdf Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Fri, 9 Sep 2022 09:45:36 +0530
Subject: [PATCH] mae/mse loss

---
 enhancer/utils/loss.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 enhancer/utils/loss.py

diff --git a/enhancer/utils/loss.py b/enhancer/utils/loss.py
new file mode 100644
index 0000000..b9c8bee
--- /dev/null
+++ b/enhancer/utils/loss.py
@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+
+
+class mean_squared_error(nn.Module):
+
+    def __init__(self,reduction="mean"):
+        super().__init__()
+
+        self.loss_fun = nn.MSELoss(reduction=reduction)
+
+    def forward(self,prediction:torch.Tensor, target: torch.Tensor):
+
+        return self.loss_fun(prediction, target)
+
+class mean_absolute_error(nn.Module):
+
+    def __init__(self,reduction="mean"):
+
+        self.loss_fun = nn.L1Loss(reduction=reduction)
+
+    def forward(self, prediction:torch.Tensor, target: torch.Tensor):
+
+        return self.loss_fun(prediction, target)
+        
+LOSS_MAP = {"mea":mean_absolute_error, "mse": mean_squared_error}
+
+