Import

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

Device’s info

print('cuda is available? ', torch.cuda.is_available())
print('device_count: ', torch.cuda.device_count())
print('current device: ', torch.cuda.current_device())
print('device name: ', torch.cuda.get_device_name(0))

Problem with CUDA version

For example, need to install corresponding versions: torch==1.2.0torchvision==0.4.0Pillow<7.0.0

pip3 install -U torch==1.2.0
pip3 install -U torchvision==0.4.0
pip3 install -U "pillow<7"