Custom autograd function list/dict input

The documentation for extending pytorch has this to say about inputting list or dicts to autograd functions:

“All kinds of Python objects are accepted here. Tensor arguments that track history (i.e., with requires_grad=True) will be converted to ones that don’t track history before the call, and their use will be registered in the graph. Note that this logic won’t traverse lists/dicts/any other data structures and will only consider Tensor s that are direct arguments to the call.”

I am interested in creating an autograd function which accepts a list of Tensors that track history which are used in the computation. The quoted documentation suggests that I would need to manually make sure that those Tensors are added to the computation graph. Is there an example of how to do this anywhere? I can’t seem to find one.


Unfortunately, this is not possible at the moment.
You can simply doing it by passing varargs: YourFunc.apply(*your_list). And when creating the function def forward(ctx, *your_list).

Hi, Do you know how I can register a tensor in the graph passed as an index argument to nn.ModuleList so that it can be trained?


I am not sure to follow, ModuleList is a torch.nn construct and unrelated to this discussion about custom Functions.

Maybe open a new topic with a code sample that shows what you’re trying to achieve?