remove class check in registry; add type assert in weight_init

2020-08-08 07:08:23 +08:00
parent f7843de45d
commit 206d9343cd
2 changed files with 5 additions and 3 deletions
--- a/model/weight_init.py
+++ b/model/weight_init.py
@@ -67,4 +67,6 @@ def generation_init_weights(module, init_type='normal', init_gain=0.02):
            # only normal distribution applies.
            normal_init(m, 1.0, init_gain)

+    assert isinstance(module, nn.Module)
    module.apply(init_func)
+