I'm not sure if this function that computes convolution uses the GPU, if not, how to modify it to use GPU?
import torch
import time
cuda = torch.device('cuda')
x = torch.randn(1,16,4,4)
t0 = time.time()
conv = torch.nn.Conv2d(16,1,(3,3))
res = conv(x)
torch.cuda.synchronize()
t1 = time.time()
print(t0, t1, t1-t0,"sec")
print(res.shape)
You can type this command at the beginning of your code:
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
If it prints "Using cuda device" means that it is using GPU
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.