测试程序运行在CPU还是GPU

import random
import time

import torch
from torch_frft.frft_module import frft

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

N = 3000
a = 0.5
# X = torch.rand(N, N, device=device)
X = [random.randint(1, 100) for _ in range(3000)]
X = torch.tensor(X)
X.to(device)
t1 = time.time()
Y1 = frft(X, a)
print(X.device)
t2 = time.time()
times = t2-t1
print(times, "秒")
print(Y1)