Hi,
From the error in the colab, you use .data[0].
.data should not be used anymore.
If the goal here is to get a python number. you should replace .data[0] by just .item() (as suggested in the error message).