refactoring __main__

aniketmaurya · Mar 17, 2024 · 5c8ad25 · 5c8ad25
1 parent d83e937
commit 5c8ad25
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/fastserve/models/__main__.py b/src/fastserve/models/__main__.py
@@ -76,6 +76,7 @@
 elif args.model == "huggingface":
     app = ServeHuggingFace(
         model_name=args.model_name,
+        use_gpu=True if args.use_gpu else False,
         device=device,
         timeout=args.timeout,
         batch_size=args.batch_size,