Fixes issues with Colab being updated to Python 3.10 and CUDA 11.8

[ada-ada-ada-art]

I changed two files in the repo: torch_utils/ops/conv2d_gradfix.py and torch_utils/ops/grid_simple_gradfix.py.

I copied the files from the StyleGAN3 repo, which has received an update to handle new PyTorch versions.

In the Colab training notebook, I also removed all the JAX and PyTorch uninstall stuff, so that this:

#Uninstall new JAX
!pip uninstall jax jaxlib -y
#GPU frontend
!pip install "jax[cuda11_cudnn805]==0.3.10" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
#CPU frontend
#!pip install jax[cpu]==0.3.10
#Downgrade Pytorch
!pip uninstall torch torchvision -y
!pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
!pip install timm==0.4.12 ftfy==6.1.1 ninja==1.10.2 opensimplex

becomes this:

!pip install timm==0.4.12 ftfy==6.1.1 ninja==1.10.2 opensimplex

[persac]

@dvschultz Would you please consider merging this PR? I can confirm that it fixes issue #45. 🙏