Как извлечь 1 срез тензора 3D-изображения?

скажем, у меня есть 3D-изображение 64x64x64. У меня также есть вектор x длиной 64.

Я хочу взять слой argmax(x) следующим образом:

2d_image = 3d_image[:,argmax(x),:]

точнее (для тензорного потока):

def extract_slice(x,3d_image):
     slice_index = tf.math.argmax(x,axis=1,output_type=tf.dtype.int32) #it has to be int for indexing
     return 3d_image[:,slice_index,:]

Ошибка:

Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got <tf.Tensor 'ArgMax_50:0' shape=(None,) dtype=int32>

np.shape параметров:

Форма 3d_image (None, 64, 64, 64, 1)

х форма (None, 64)

Форма slice_index (None,)

-> измерение 1 в форме 3d_image связано с тем, что это образец из массива. Я не думаю, что это имеет значение

Я знаю, что форма None - это размер партии, который неизвестен, но другие выглядят отлично... так в чем проблема?

Насколько я понимаю, похоже, что индекс не int32, но я действительно привел его к tf.int, так в чем может быть проблема?? может быть, int32 отличается от tf.int32? или, может быть, метод индекса, который я использовал, недействителен в тензорном потоке? может быть, это должна быть функция что-то вроде этого: tf.index(изображение,[:,slice_index,:])..?

Спасибо!