first commit

utils/losses.py

@@ -0,0 +1,15 @@
+from typing import Optional, Callable
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def get_loss_fn(loss_name: str,
+                delta: Optional[float] = 1.0,
+                beta: Optional[float] = 1.0) -> Callable:
+    return {'mse': F.mse_loss,
+            'mae': F.l1_loss,
+            'huber': partial(F.huber_loss, delta=delta),
+            'smooth_l1': partial(F.smooth_l1_loss, beta=beta)}[loss_name]