From 195d12fdff7140392b513ead655df500d3ba25b6 Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Wed, 10 Apr 2024 09:34:34 -0700
Subject: [PATCH] Add missing device argument to inference.py dataloader

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

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,
     )