[Bugfix] Correctly call cudaProfilerStop in benchmarks script (#14183)

Signed-off-by: Brayden Zhong <b8zhong@uwaterloo.ca>
This commit is contained in:
Brayden Zhong
2025-03-06 19:42:49 -05:00
committed by GitHub
parent ad60bbb2b2
commit c34eeec58d
6 changed files with 5 additions and 6 deletions

View File

@@ -40,7 +40,7 @@ def main(num_tokens: int,
end_time = time.perf_counter()
if profile:
torch.cuda.cudart().cudaProfilerStart()
torch.cuda.cudart().cudaProfilerStop()
return (end_time - start_time) / num_iters
# Warmup.