Careful: if you only consider the first variable found in an iterable or mapping, the backward hook won't be registered for the rest!

Also, you might want to use the collections interface to determine whether something is list-like or dict-like (example).

Thanks for pointing out the backward hook issue. So the original code iterated through the result picking up the first as it went deeper. What I want to do is enable named outputs to be returned from forward calls (in a dict).

So would it be safe to assume that the first returned item from forward will always be the output and hence all its Variables should be hooked? That seems to be implicit in the original code.

that seems to be implicit in the original code

I don't think the code was designed with dict inputs/outputs in mind. If you check the list of PRs, you'll notice one for enabling dict inputs to DataParallel. For that one, the original code actually assumed, in several places, that there is always one positional input.

Since it seems like there's a recent community push for named variables, I'd say make whatever changes are necessary to ensure the expected behavior (e.g., hooking all outputs) .

We don't need all outputs to register the hooks properly, one is enough. However this PR won't really fix it. If you're going to return such dict: {'a': output, 'b': 4}, then in different interpreter runs the hash seeds will be different, and the first element of the .values() will be either 4 or output, so the code will randomly crash from time to time. I think we'll just need to add some additional functionality to autograd so we can track the outputs properly, and remove any constraints on the inputs and outputs to the modules. But that's out of scope for this PR.

@apaszke Valid point - I agree that we've a lot of implicit assumptions here that any made up rule on what should be returned where is going to break.

Happy to help on the autograd suggestion and close this PR - where shall I start? If you give me a few pointers I can look into it.

Today I ran into the same issue that this PR tries to fix, and unknowingly submitted on my own PR in #1239. I find the suggested change quite useful, and consider it an improvement to the current code which allows (single) variables or lists of variables and picks out the first item for hooking if list.

@apaszke

We don't need all outputs to register the hooks properly, one is enough.

Can you clarify why this is true?

However this PR won't really fix it. If you're going to return such dict: {'a': output, 'b': 4}, then in different interpreter runs the hash seeds will be different, and the first element of the .values() will be either 4 or output, so the code will randomly crash from time to time.

If module.forward() returns a list or a dict, can we not assume that all items are variables? If not, is it better to filter non-variable items and hook the rest, instead of hooking just the first element. (Again, I'm probably missing why the first element is enough.)

We'll need to figure out a solution to this problem and send a PR ourselves. I'm closing this for now, thanks!