Olá, Victor!
Para que eu possa analisar melhor o que pode estar acontecendo no seu código, poderia, por favor, compartilhar o link do seu Colab? Assim consigo visualizar sua implementação e te ajudar de forma mais assertiva.
Fico no aguardo e a disposição :)