feat: add repr for LinearTransform

2022-04-01 10:13:25 +02:00
parent 46ff1c4eb1
commit 85f75bb28c
1 changed files with 3 additions and 0 deletions
--- a/prototorch/core/transforms.py
+++ b/prototorch/core/transforms.py
@@ -38,6 +38,9 @@ class LinearTransform(torch.nn.Module):
    def forward(self, x):
        return x @ self._weights

+    def extra_repr(self):
+        return f"weights: (shape: {tuple(self._weights.shape)})"
+

 # Aliases
 Omega = LinearTransform