fixed bug on replicate

2022-07-15 21:21:16 -04:00
parent cde7900adc
commit 00d9363172
2 changed files with 6 additions and 5 deletions
@@ -11,7 +11,8 @@ class ReplicatePredictor(BasePredictor):
        self.model = MinDalle(
            is_mega=True, 
            is_reusable=True, 
-            dtype=torch.float32
+            dtype=torch.float32,
+            device='cuda'
        )

    def predict(