Update main.py

main.py

@@ -14,8 +14,8 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# device = "cpu"
-device = "cuda"
+# device = torch.device("cpu")
+device = torch.device("cuda")
 
 # comment in/out the models you want to use
 # RAM requirements: ~16GB x #models (+ ~4GB overhead)