diff options
| author | David Riazati <davidriazati@fb.com> | 2019-04-11 15:33:51 -0700 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2019-04-11 15:38:06 -0700 |
| commit | fea0a0be5375e9b09c7048dd0f4ec225b17cf225 (patch) | |
| tree | 624e1f1b2678f95b09c57397a904bacd7283fb18 /torch | |
| parent | 4ae59e47440429fdec40d728facfc8fb3730c16e (diff) | |
| download | pytorch-fea0a0be5375e9b09c7048dd0f4ec225b17cf225.tar.gz pytorch-fea0a0be5375e9b09c7048dd0f4ec225b17cf225.tar.bz2 pytorch-fea0a0be5375e9b09c7048dd0f4ec225b17cf225.zip | |
Support attributes when copying modules (#19040)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/19040
ghimport-source-id: 37933efd717795751283cae8141e2e2caaae2e95
Reviewed By: eellison
Differential Revision: D14895573
Pulled By: driazati
fbshipit-source-id: bc2723212384ffa673d2a8df2bb57f38c62cc104
Diffstat (limited to 'torch')
| -rw-r--r-- | torch/csrc/jit/script/module.h | 10 |
1 files changed, 6 insertions, 4 deletions
diff --git a/torch/csrc/jit/script/module.h b/torch/csrc/jit/script/module.h index e5b3a65823..ef4fd25076 100644 --- a/torch/csrc/jit/script/module.h +++ b/torch/csrc/jit/script/module.h @@ -425,11 +425,13 @@ struct TORCH_API Module { parameter_remap[param] = curr->parameter_slot(param.name()); } for (auto& attr : get_attributes()) { - if (!attr.type()->isSubtypeOf(TensorType::get())) { - continue; + if (attr.type()->isSubtypeOf(TensorType::get())) { + curr->register_buffer(attr.name(), attr.value().toTensor()); + parameter_remap[attr] = *curr->find_buffer(attr.name()); + } else { + curr->register_attribute(attr.name(), attr.type(), attr.value()); + parameter_remap[attr] = *curr->find_attribute(attr.name()); } - curr->register_buffer(attr.name(), attr.value().toTensor()); - parameter_remap[attr] = *curr->find_buffer(attr.name()); } for (auto& mod : get_modules()) { names.push_back(mod->name()); |