Skip to content

Commit 65f56f4

Browse files
authored
[example] gpt demo more accuracy tflops (#2178)
1 parent ab54fed commit 65f56f4

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

examples/language/gpt/train_gpt_demo.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,7 @@ def main():
283283
optimizer.sync_grad()
284284
optimizer.step()
285285
logger.info(get_mem_info(prefix=f'[{n+1}/{NUM_STEPS}] Optimizer step '), ranks=[0])
286+
torch.cuda.synchronize()
286287
step_time = time() - start
287288
logger.info(
288289
f'[{n+1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}',

0 commit comments

Comments
 (0)