From cabd26d018e29ee86c6ea6ae0c771ec8b9b5867e Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Wed, 14 May 2025 13:51:43 -0700
Subject: [PATCH] Fix missing device arg in train.py, fix #2488

---
 train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/train.py b/train.py
index fda3c184..11b783e8 100755
--- 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,