cuda

1 year ago · aa25e65bad
parent 37d63776c9
commit aa25e65bad
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