pytorch-GPU检测代码

测试GPU

import torch
flag = torch.cuda.is_available()
print(flag)
ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda()) 

结果：
True
cuda:0
NVIDIA GeForce RTX 2080 Ti
tensor([[0.0535, 0.1229, 0.6272],
[0.3767, 0.1368, 0.9835],
[0.7672, 0.5263, 0.2509]], device=‘cuda:0’)

import math, random
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.autograd as autograd
import torch.nn.functional as F
from collections import  Counter
from collections import deque
import matplotlib.pyplot as plt