May I ask what version do you use for torch, torch_vision and Cuda? It seems to have some problems running your code on colab because of the version.