[FEATURE] Remove utility modules and add wrappers instead

prototorch/modules/utils.py

@@ -1,16 +0,0 @@
-"""ProtoTorch utility modules."""
-
-import torch
-
-
-class LambdaLayer(torch.nn.Module):
-    def __init__(self, fn, name=None):
-        super().__init__()
-        self.fn = fn
-        self.name = name or fn.__name__ # lambda fns get <lambda>
-
-    def forward(self, *args, **kwargs):
-        return self.fn(*args, **kwargs)
-
-    def extra_repr(self):
-        return self.name

prototorch/modules/wrappers.py

@@ -0,0 +1,36 @@
+"""ProtoTorch Wrappers."""
+
+import torch
+
+
+class LambdaLayer(torch.nn.Module):
+    def __init__(self, fn, name=None):
+        super().__init__()
+        self.fn = fn
+        self.name = name or fn.__name__  # lambda fns get <lambda>
+
+    def forward(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+    def extra_repr(self):
+        return self.name
+
+
+class LossLayer(torch.nn.modules.loss._Loss):
+    def __init__(self,
+                 fn,
+                 name=None,
+                 size_average=None,
+                 reduce=None,
+                 reduction: str = "mean") -> None:
+        super().__init__(size_average=size_average,
+                         reduce=reduce,
+                         reduction=reduction)
+        self.fn = fn
+        self.name = name or fn.__name__  # lambda fns get <lambda>
+
+    def forward(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+    def extra_repr(self):
+        return self.name