From db82cd9f35bfd938e66433b836ed51b5afb76b2a Mon Sep 17 00:00:00 2001 From: Gene Date: Sun, 3 Nov 2019 14:24:28 +0800 Subject: [PATCH] Update pruning_engine.py set momentum to zero for pruned neurons (matching the description) --- pruning_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pruning_engine.py b/pruning_engine.py index d1497cc..1e5f18f 100755 --- a/pruning_engine.py +++ b/pruning_engine.py @@ -900,7 +900,7 @@ def set_momentum_zero_sgd(self, optimizer=None): if not self.prune_layers[layer]: continue for unit in range(len(self.pruning_gates[layer])): - if not self.pruning_gates[layer][unit]: + if self.pruning_gates[layer][unit]: continue if 'momentum_buffer' in optimizer.state[self.parameters[layer]].keys(): optimizer.state[self.parameters[layer]]['momentum_buffer'][unit] *= 0.0