python - How can use index in tensorflow?

Question

I have three 2-dim tensor in the training. I made like this below, but I get an error

Cannot convert an unknown Dimension to a Tensor: ?

mu = params[:, :self.num_mixture]
sigma = params[:, self.num_mixture:self.num_mixture * 2]
alpha = params[:, self.num_mixture * 2:self.num_mixture * 3]
max_alpha = tf.cast(tf.argmax(alpha, axis=1 ), 'int32')

mu = tf.gather_nd(mu,tf.stack([tf.range(max_alpha.shape[0])[...,tf.newaxis], mu[...,tf.newaxis]], axis=2))
sigma = tf.gather_nd(sigma,tf.stack([tf.range(max_alpha.shape[0])[...,tf.newaxis], sigma[...,tf.newaxis]], axis=2))

How can I fix it?