如何在 torch.nn.Conv2d() function 中使用 GPU

Question

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) 

Answer 1

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