就三点:
- 申明用GPU
- 把model放在GPU上
- 把数据和标签放在GPU上
主要介绍第一步申明GPU
两种方法
4. device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
默认是第一张卡
5. 指定具体某张或者多张卡
import os
os.environ[“CUDA_VISIBLE_DEVICES”]=“2,3”
import torch
就三点:
主要介绍第一步申明GPU
两种方法
4. device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
默认是第一张卡
5. 指定具体某张或者多张卡
import os
os.environ[“CUDA_VISIBLE_DEVICES”]=“2,3”
import torch