add a hook for on_tng_metrics so that users get access to the grad_norm and mem_map dicts.

This commit is contained in:
Cinjon Resnick
2019-07-16 12:51:48 -04:00
parent 28cfddbe65
commit fbd3873a0f
2 changed files with 9 additions and 5 deletions
+3
View File
@@ -19,3 +19,6 @@ class ModelHooks(torch.nn.Module):
def on_post_performance_check(self):
pass
def on_tng_metrics(self, metrics):
pass