cuda

Former-commit-id: 52772820c70c995cc1b3348323fc500aa69a3fe5
2 years ago · f2f702f1fb
parent ee60d8d591
commit f2f702f1fb
1 changed files with 1 additions and 1 deletions
--- a/swarms/models/fuyu.py
+++ b/swarms/models/fuyu.py
@ -33,7 +33,7 @@ class Fuyu:
    def __init__(
        self,
        pretrained_path: str = "adept/fuyu-8b",
-        device_map: str = "cuda",
+        device_map: str = "cuda:0",
        max_new_tokens: int = 7,
    ):
        self.pretrained_path = pretrained_path