From 6841ada9bcd228f4a24c9c7909cec2e85c4bc0df Mon Sep 17 00:00:00 2001 From: Suharsh Sivakumar Date: Mon, 12 Feb 2018 12:50:07 -0800 Subject: [PATCH] Also add quantization step node to MODEL_VARIABLES collection. PiperOrigin-RevId: 185420228 --- tensorflow/contrib/quantize/python/common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/contrib/quantize/python/common.py b/tensorflow/contrib/quantize/python/common.py index 3a1fa61..9e76549 100644 --- a/tensorflow/contrib/quantize/python/common.py +++ b/tensorflow/contrib/quantize/python/common.py @@ -114,7 +114,9 @@ def CreateOrGetQuantizationStep(): dtype=dtypes.int64, initializer=init_ops.zeros_initializer(), trainable=False, - collections=[ops.GraphKeys.GLOBAL_VARIABLES]) + collections=[ + ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.MODEL_VARIABLES + ]) with g.name_scope(quantization_step_tensor.op.name + '/'): # We return the incremented variable tensor. Since this is used in conds # for quant_delay and freeze_bn_delay, it will run once per graph -- 2.7.4