Support for Nvidia Ampere cards

[entmike]

Just an FYI, I was playing around with this repo all morning (awesome work!) and was wondering why it was so slow. It turns out that it was using CPU and not CUDA/GPU. (It would be nice if you spammed a warning for others in case they are not as stubborn as me) I am using a newer Ampere card that needs CUDA 11, so I had to solve the issue as shown below. I hope this might help others. Note that this is similar to another closed issue by somebody else, however I had to make further changes such as adding cudatoolkit=11.1 to mine.

name: pytorch-nst
channels:
  - defaults
  - pytorch
  - conda-forge
dependencies:
  - cudatoolkit=11.1
  - python=3.7.6
  - pip=20.0.2
  - matplotlib=3.1.3
  - torchvision
  - torchaudio
  - pip:
    - numpy==1.18.1
    - opencv-python==4.2.0.32
    - conda-forge::pytorch

[Finerrkekz]

Managed to get it to work by running this line in anaconda

pip install torch===1.7.1+cu110 torchvision===0.8.2+cu110 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html