Skip to content

Commit

Permalink
Handle None in modules list.
Browse files Browse the repository at this point in the history
It's often useful to add None to an nn.ModuleList to keep the indexing
of the module list to match some other property.
  • Loading branch information
colesbury authored and soumith committed Jul 3, 2017
1 parent 39edc37 commit b4414c0
Showing 1 changed file with 9 additions and 4 deletions.
13 changes: 9 additions & 4 deletions torch/nn/parallel/replicate.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,15 @@ def replicate(network, devices):

for i, module in enumerate(modules):
for key, child in module._modules.items():
module_idx = module_indices[child]
for j in range(num_replicas):
replica = module_copies[j][i]
replica._modules[key] = module_copies[j][module_idx]
if child is None:
for j in range(num_replicas):
replica = module_copies[j][i]
replica._modules[key] = None
else:
module_idx = module_indices[child]
for j in range(num_replicas):
replica = module_copies[j][i]
replica._modules[key] = module_copies[j][module_idx]
for key, param in module._parameters.items():
if param is None:
for j in range(num_replicas):
Expand Down

0 comments on commit b4414c0

Please sign in to comment.