[Fix] fix eos trim inconsistency (#1650)

This commit is contained in:
Ying Sheng
2024-10-13 01:07:09 -07:00
committed by GitHub
parent c3f2fc5a7a
commit 4876117171
7 changed files with 77 additions and 27 deletions

View File

@@ -690,3 +690,10 @@ def pytorch_profile(name, func, *args, data_size=-1):
prof.export_chrome_trace(f"trace/{name}_{step_counter}.json")
step_counter += 1
return result
def first_rank_print(*args, **kwargs):
if torch.cuda.current_device() == 0:
print(*args, **kwargs)
else:
pass