nemo.utils.get_summary returns a dictionary + prettyprint information…

… now
FrancescoConti · Mar 26, 2021 · cd103c2 · cd103c2
1 parent 75e91b9
commit cd103c2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/utils.py b/nemo/utils.py
@@ -252,7 +252,7 @@ def hook(module, input, output):
     s += "----------------------------------------------------------------" + "\n"
     if verbose:
         logging.info(s)
-    return summary, s
+    return { 'dict': summary, 'prettyprint': s, 'biggest_activation': output_size, 'params_size': params_size }
 
 def get_intermediate_activations(net, test_fn, *test_args, **test_kwargs):
     l = len(list(net.named_modules()))