tf2 - from colab to local machine - problems with TimeDistributed

[NGluna03]

Hi to everyone!
I'm trying to apply this code in my own dataset. I setted it up on google colab using tensorflow v2 without problems but now I need to run it in my local machine. I tryed to install the same packages versions getting them through the command "pip list -v".
However, I'm running into the following error:
"ValueError: Dimension size must be evenly divisible by 18944 but is 150528 for 'time_distributed_2/Reshape' (op: 'Reshape') with input shapes: [6,7,7,512], [3] and with input tensors computed as partial shapes: input[1] = [?,37,512]."

I'm using a feature map of 37 x 37. Considering that 37 x 512 = 18944, it seams that the first TimeDistributed (inside function "classifier_layer"), is not understanding the correct input at the following line:

out = TimeDistributed(Flatten(name='flatten'))(out_roi_pool)

Did someone run into this problem?
Thank you in advance for your reply.