From 65f56f49e8ff5db4fe9b1911f90f8d1950fe1c47 Mon Sep 17 00:00:00 2001 From: Jiarui Fang Date: Thu, 22 Dec 2022 20:51:35 +0800 Subject: [PATCH] [example] gpt demo more accuracy tflops (#2178) --- examples/language/gpt/train_gpt_demo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/language/gpt/train_gpt_demo.py b/examples/language/gpt/train_gpt_demo.py index 4db9d66e4..3b22f05a6 100644 --- a/examples/language/gpt/train_gpt_demo.py +++ b/examples/language/gpt/train_gpt_demo.py @@ -283,6 +283,7 @@ def main(): optimizer.sync_grad() optimizer.step() logger.info(get_mem_info(prefix=f'[{n+1}/{NUM_STEPS}] Optimizer step '), ranks=[0]) + torch.cuda.synchronize() step_time = time() - start logger.info( f'[{n+1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}',