From 2e0d147fd776c766dc9b24a0b68b08809a907af7 Mon Sep 17 00:00:00 2001
From: Guy Jacob <guy.jacob@intel.com>
Date: Sun, 6 Oct 2019 12:44:54 +0300
Subject: [PATCH] Quantizer: Re-move model to device at the end of
prepare_model
---
distiller/quantization/quantizer.py | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/distiller/quantization/quantizer.py b/distiller/quantization/quantizer.py
index 369271c..2b3ac78 100644
--- a/distiller/quantization/quantizer.py
+++ b/distiller/quantization/quantizer.py
@@ -220,6 +220,8 @@ class Quantizer(object):
summary_graph = distiller.SummaryGraph(self.model, dummy_input)
self.adjacency_map = summary_graph.adjacency_map(dedicated_modules_only=False)
+ model_device = distiller.model_device(self.model)
+
self._pre_prepare_model(dummy_input)
self._pre_process_container(self.model)
@@ -247,6 +249,9 @@ class Quantizer(object):
self._post_prepare_model()
+ # Re-transfer model to the device it was on, in case the quantizer created new parameters/buffers
+ self.model.to(model_device)
+
msglogger.info('Quantized model:\n\n{0}\n'.format(self.model))
def _pre_prepare_model(self, dummy_input):
--
GitLab