Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
torch.cuda.synchronize()
torch.accelerator.synchronize