I replied on the GitHub issue. I couldn’t reproduce it on Colab. Can you try setting up a fresh environment?