diff --git a/inference.py b/inference.py
index f19dba28..d8cb3dee 100755
--- a/inference.py
+++ b/inference.py
@@ -241,6 +241,7 @@ def main():
         batch_size=args.batch_size,
         use_prefetcher=True,
         num_workers=workers,
+        device=device,
         **data_config,
     )