from tensorflow.python.ops import array_ops

      loss_vec = array_ops.reshape(loss_vec, shape=(-1,))
      loss_vec = math_ops.mul(
          loss_vec, array_ops.reshape(weight_tensor, shape=(-1,)))
      return math_ops.div(
