Fix missing device arg in train.py, fix #2488

2025-06-03 15:01:08 +08:00 · 2025-05-14 13:51:43 -07:00 · 2025-05-14 13:51:43 -07:00 · cabd26d018
commit cabd26d018
parent 6b302f27a3
1 changed files with 1 additions and 0 deletions
--- a/train.py
+++ b/train.py
@ -892,6 +892,7 @@ def main():
                optimizer,
                train_loss_fn,
                args,
+                device=device,
                lr_scheduler=lr_scheduler,
                saver=saver,
                output_dir=output_dir,