Bug in spatial_transformer_network's code

[AyushSharma2000]

Environment:

Google Colaboratory.
Tensorflow version 2.5.0
Keras version 2.4.x

Tried running the train.py file in examples/Spatial_Transformer_Networks/ and received an error on passing the MNIST dataset present in the STN repository. No changes were made to the code other than changing the file path.

The second screenshot shows the remaining portion of the last line of error present in the first screenshot.

[zjhJOJO]

i met the same bug, can anyone helpful offer a solution?

[Ataei67]

change the compute_output_shape function in layer.py file to this solved my problem.
def compute_output_shape(self, input_shapes): height, width = self.output_size num_channels = input_shapes[0][-1] return tf.constant([1, height, width, num_channels])