InvalidArgumentError: PartialTensorShape: несовместимые фигуры во время слияния: [1,2] против [2,2]

Question

Информация о системе
-OS Платформа и распространение (например, Linux Ubuntu 16.04): Windows 10 64 бит

-Мобильное устройство (например, iPhone 8 , Pixel 2, Samsung Galaxy), если проблема возникает на мобильном устройстве: No

-TensorFlow установлен из (источника или двоичного файла): NVIDIA (https://docs.nvidia.com/deeplearning/frameworks/tensorflow-release-notes)

-TensorFlow версия (используйте команду ниже): 2.1

- Python версия: 3.6.9

-Количество графических процессоров: 1

Здравствуйте, я использую свою модель NN и он прекрасно работает в некоторые эпохи, и внезапно он ломается, показывая следующую проблему:

_ InvalidArgumentError : PartialTensorShape: несовместимые фигуры во время слияния: [1,2] против [2,2] [[node TensorArrayV2Stack / TensorListStack (определено в /workspace/code/model_NN.py:268)]] [Op: __ forward_p2d_fluoro_prediction_117955]

Ошибки могут возникать в результате операции ввода. Операции источника ввода, связанные с узлом TensorArrayV2Stack / TensorListStack: while (определено в /workspace/code/model_NN.py:261)

Стек вызовов функций: p2d_fluoro_prediction_

Для меня это странное поведение потому что каждый раз он ломается после другой эпохи. Функция, вызывающая проблему, следующая:

@tf.function    
    def p2d_fluoro_prediction(self,heat_map):
        points = tf.TensorArray(tf.int64, size=self.nb_points,dynamic_size=True)
        for i in tf.range(self.nb_points):
            hm = heat_map[:,:,i]
            heat_max = tf.reduce_max(hm)
            point = tf.where(hm[:,:]==heat_max)
            point = tf.reverse(point, [1])
            points = points.write(i,point)          
        points = points.stack()    
        return tf.cast(tf.squeeze(points,axis = 1),dtype ='float32')

Я надеюсь, что вы можете помочь мне с этой проблемой.