Skip to content

Instantly share code, notes, and snippets.

@mypapit

mypapit/install-pytorch-cuda.py

Created August 10, 2024 16:21
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save mypapit/7f9e286ae77be90404f9c2a2b3240a8b to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save mypapit/7f9e286ae77be90404f9c2a2b3240a8b to your computer and use it in GitHub Desktop.
Download ZIP
Raw
install-pytorch-cuda.py
#conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia.
import torch
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# Assuming that we are on a CUDA machine, this should print a CUDA device:
print(device)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment