Add weight_mask to calc_mask kwargs parameter (#2024)

microsoft · Feb 10, 2020 · 8092c8b · 8092c8b
1 parent b210695
commit 8092c8b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/sdk/pynni/nni/compression/torch/compressor.py b/src/sdk/pynni/nni/compression/torch/compressor.py
@@ -225,15 +225,17 @@ def __init__(self, module, module_name, module_type, config, pruner):
         # config and pruner
         self.config = config
         self.pruner = pruner
+        self.registered_buffers = {}
 
         # register buffer for mask
         self.register_buffer("weight_mask", torch.ones(self.module.weight.shape))
+        self.registered_buffers['weight_mask'] = self.weight_mask
         if hasattr(self.module, 'bias') and self.module.bias is not None:
             self.register_buffer("bias_mask", torch.ones(self.module.bias.shape))
         else:
             self.register_buffer("bias_mask", None)
+        self.registered_buffers['bias_mask'] = self.bias_mask
         # register user specified buffer
-        self.registered_buffers = {}
         for name in self.pruner.buffers:
             self.register_buffer(name, self.pruner.buffers[name].clone())
             self.registered_buffers[name] = getattr(self, name)