Unverified Commit 65f56f49 authored by Jiarui Fang's avatar Jiarui Fang Committed by GitHub
Browse files

[example] gpt demo more accuracy tflops (#2178)

parent ab54fed2
......@@ -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}',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment