[FIXED] pytorch hook function is not executed

Issue

May I know why the pytorch hook function does not work ?

Solution

You might want to use torch.nn.Module.named_modules instead of torch.nn.Module.named_children. The latter will only return immediate child modules. In your case graph's immediate child is cells, so you won't be looping over modules inside of cells, i.e. the layers defined inside the ModuleList.

Either use named_modules:

for name, module in graph.named_modules()
    pass

Or use named_children on graph.cells directly:

for name, module in graph.cells.named_children():
    pass

However, the latter alternative won't scale if you ever decide to add additional child modules to Graph.

Answered By - Ivan