import torch
time_pool = 0
for input_tensor,label in dataloader:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record(stream=torch.cuda.current_stream())
output_tensor = model(input_tensor)
end.record(stream=torch.cuda.current_stream())
end.synchronize()
time_pool += start.elapsed_time(end)
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)