From cb3049bd69c23b4f895c55caabb54ac2d66d5586 Mon Sep 17 00:00:00 2001
From: Neta Zmora <neta.zmora@intel.com>
Date: Tue, 6 Nov 2018 22:54:49 +0200
Subject: [PATCH] Bug fix: Thinning API was not updated after changing the
scheduler callback signature
We recently changed the signature of the on_minibatch_begin() callback
from the scheduler (added 'meta') and the callback client
in the thinning module was not updated.
---
distiller/thinning.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/distiller/thinning.py b/distiller/thinning.py
index 2a89142..ab6cff1 100755
--- a/distiller/thinning.py
+++ b/distiller/thinning.py
@@ -386,7 +386,7 @@ class FilterRemover(ScheduledTrainingPolicy):
self.thinning_func(model, zeros_mask_dict, self.arch, self.dataset, optimizer=optimizer)
self.done = True
- def on_minibatch_begin(self, model, epoch, minibatch_id, minibatches_per_epoch, zeros_mask_dict, optimizer):
+ def on_minibatch_begin(self, model, epoch, minibatch_id, minibatches_per_epoch, zeros_mask_dict, meta, optimizer):
# We hook onto the on_minibatch_begin because we want to run after the pruner which sparsified
# the tensors. Pruners configure their pruning mask in on_epoch_begin, but apply the mask
# only in on_minibatch_begin
--
GitLab