Skip to content
Snippets Groups Projects
Commit e5b14d4b authored by André Anjos's avatar André Anjos :speech_balloon:
Browse files

[utils.summary] Fix docformatting with small hack

parent eb3ae3d2
No related branches found
No related tags found
No related merge requests found
Pipeline #69264 passed
......@@ -27,8 +27,8 @@ def summary(model):
nparam : int
number of parameters
"""
def repr(model):
# ignore this space!
def _repr(model):
# We treat the extra repr like the sub-module, one item per line
extra_lines = []
extra_repr = model.extra_repr()
......@@ -38,13 +38,13 @@ def summary(model):
child_lines = []
total_params = 0
for key, module in model._modules.items():
mod_str, num_params = repr(module)
mod_str, num_params = _repr(module)
mod_str = _addindent(mod_str, 2)
child_lines.append("(" + key + "): " + mod_str)
total_params += num_params
lines = extra_lines + child_lines
for name, p in model._parameters.items():
for _, p in model._parameters.items():
if hasattr(p, "dtype"):
total_params += reduce(lambda x, y: x * y, p.shape)
......@@ -60,4 +60,4 @@ def summary(model):
main_str += f", {total_params:,} params"
return main_str, total_params
return repr(model)
return _repr(model)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment