Is this the right way to time profile the model prediction on a test sample?
import torch.autograd.profiler as profiler
x.to("cuda:0") # test sample
model.eval()
with profiler.profile(use_cuda=True) as prof:
    with torch.no_grad():
        pred = model(x)
print(prof)