Update how the model is printed

2021-06-07 16:58:00 +02:00 · 2021-06-07 16:58:00 +02:00 · 34ffeb95bc
commit 34ffeb95bc
parent 3aa33fd182
1 changed files with 4 additions and 2 deletions
--- a/prototorch/models/abstract.py
+++ b/prototorch/models/abstract.py
@ -12,8 +12,10 @@ from prototorch.modules import WTAC, LambdaLayer

 class ProtoTorchBolt(pl.LightningModule):
    def __repr__(self):
-        super_repr = super().__repr__()
-        return f"ProtoTorch Bolt:\n{super_repr}"
+        surep = super().__repr__()
+        indented = "".join([f"\t{line}\n" for line in surep.splitlines()])
+        wrapped = f"ProtoTorch Bolt(\n{indented})"
+        return wrapped


 class PrototypeModel(ProtoTorchBolt):