Saving from colab to google drive with pytorch makes black image

Question

I mounted a google drive with Colab:

from google.colab import drive
drive.mount(\'/content/drive\')

this works fine. Using pytorch to save a