fix remove_weight_norm() with wrong block (#36)

seungwonpark · Dec 3, 2019 · a7c2351 · a7c2351
1 parent 00cb82d
commit a7c2351
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model/res_stack.py b/model/res_stack.py
@@ -31,6 +31,6 @@ def forward(self, x):
 
     def remove_weight_norm(self):
         for block, shortcut in zip(self.blocks, self.shortcuts):
-            nn.utils.remove_weight_norm(block[1])
-            nn.utils.remove_weight_norm(block[3])
+            nn.utils.remove_weight_norm(block[2])
+            nn.utils.remove_weight_norm(block[4])
             nn.utils.remove_weight_norm(shortcut)