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

rwightman · rwightman · commit cabd26d018e2 · 2025-05-14T13:51:43.000-07:00
diff --git 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,
