fix: add missing device argument in image_demo

2024-09-11 10:36:38 +08:00 · 2024-09-11 10:36:38 +08:00 · 63b563545b
parent 17a886cb58
commit 63b563545b
1 changed files with 1 additions and 1 deletions
--- a/demo/image_demo.py
+++ b/demo/image_demo.py
@ -27,7 +27,7 @@ def main():
    try:
        pretrained = args.checkpoint or True
        inferencer = ImageClassificationInferencer(
-            args.model, pretrained=pretrained)
+            args.model, device=args.device, pretrained=pretrained)
    except ValueError:
        raise ValueError(
            f'Unavailable model "{args.model}", you can specify find a model '