cuda

preprocess.py

@@ -19,9 +19,9 @@ class Preprocessor:
             return
         elif name == "NormalBae":
             print("Loading NormalBae")
-            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cpu")
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
             torch.cuda.empty_cache()
-            gc.collect()
+            # gc.collect()
         # elif name == "Canny":
         #     self.model = CannyDetector()
         else: