Fix pruner base class bug. (IntelLabs#131)

caihengyu520 · Jan 22, 2019 · 994e58d · 994e58d
1 parent 2438116
commit 994e58d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distiller/pruning/pruner.py b/distiller/pruning/pruner.py
@@ -36,5 +36,5 @@ def threshold_model(model, threshold):
  """
  for name, p in model.named_parameters():
  if 'weight' in name:
- mask = distiller.threshold_mask(param.data, threshold)
+ mask = distiller.threshold_mask(p.data, threshold)
  p.data = p.data.mul_(mask)