Iterable of layers to cuda

Suppose that I wanted to keep a dictionary of layers in my model object, when I call model.cuda(), any layer that I have defined as an attribute will be automatically converted to its cuda counterpart, however none of the layers that I have inside of an iterable (list, dict, etc;) are affected by the call. Has anyone ever wanted this functionality, and if so, what would be a good clever way of converting the objects inside of the iterables to cuda?

ModuleList (http://pytorch.org/docs/master/nn.html#modulelist) is designed to be a list-like container that also supports .cuda() and other methods. There isn’t a similar one for dict yet. Some discussions on ModuleDict are at https://github.com/pytorch/pytorch/issues/4048. There is supposed to be a PR soon according to the thread.

thanks, it seems like a very useful addition and I may end up working on it if ModuleList doesn’t accomplish what I need and no progress has been made.

Now that I that I think about it, this would be very useful in the case of multitask learning. Say you had a ModuleDict that contained multiple output layers, each with a key for the particular task. You could then apply a loss based upon groups of key values, removing the need to infer the position in a list. Hmmm…

1 Like