import tensorflow as tf

        grads_pruned += [(grad[0] * self.masks[idx_mask], grad[1])]

    return grads_pruned

  def __save_model(self):
    """Save the current model."""

    save_path = self.saver_train.save(self.sess_train, FLAGS.ws_save_path, self.global_step)
    tf.logging.info('model saved to ' + save_path)

  def __restore_model(self, is_train):
    """Restore a model from the latest checkpoint files.
    Args:
    * is_train: whether to restore a model for training
    """

