julius/prototorch_models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update how the model is printed

Browse Source
This commit is contained in:
Jensun Ravichandran 2021-06-07 16:58:00 +02:00
parent 3aa33fd182
commit 34ffeb95bc
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
prototorch/models/abstract.py
View File

@ -12,8 +12,10 @@ from prototorch.modules import WTAC, LambdaLayer
class ProtoTorchBolt(pl.LightningModule): class ProtoTorchBolt(pl.LightningModule):
def __repr__(self): def __repr__(self):
super_repr = super().__repr__() surep = super().__repr__()
return f"ProtoTorch Bolt:\n{super_repr}" indented = "".join([f"\t{line}\n" for line in surep.splitlines()])
wrapped = f"ProtoTorch Bolt(\n{indented})"
return wrapped
class PrototypeModel(ProtoTorchBolt): class PrototypeModel(ProtoTorchBolt):