frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-26 16:44:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

C++ unregister_module function for Module (#26088)

Browse Source 
Summary:
This PR adds ```unregister_module``` to ```nn::Module``` and ```erase``` function to ```OrderedDict```.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/26088

Differential Revision: D17360058

Pulled By: yf225

fbshipit-source-id: f1f375b4751317da85b8da1458e092fe2405ceec
This commit is contained in:
Shahriar
2019-09-12 18:23:25 -07:00
committed by Facebook Github Bot
parent 1d87090051
commit 18a0040fec
5 changed files with 58 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
test/cpp/api/module.cpp
View File

@ -115,6 +115,17 @@ TEST_F(ModuleTest, ReplaceModule) {
ASSERT_EQ(model->l1.get(), model->named_modules()["l1"]->as<Linear>());
}
TEST_F(ModuleTest, UnregisterModule) {
struct TestModel : public torch::nn::Module {};
TestModel model;
ASSERT_THROWS_WITH(
model.unregister_module("linear"),
"No Module with name `linear` is registered");
model.register_module("linear", torch::nn::Linear(3, 4));
model.unregister_module("linear");
ASSERT_TRUE(model.children().empty());
}
TEST_F(ModuleTest, RegisterParameterThrowsForEmptyOrDottedName) {
struct TestModel : public torch::nn::Module {};
ASSERT_THROWS_WITH(