Update src/chatterbox/mtl_tts.py

src/chatterbox/mtl_tts.py

@@ -162,9 +162,14 @@ class ChatterboxMultilingualTTS:
         ckpt_dir = Path(ckpt_dir)
 
         ve = VoiceEncoder()
-        ve.load_state_dict(
-            torch.load(ckpt_dir / "ve.pt", weights_only=True)
-        )
+        # ve.load_state_dict(
+        #     torch.load(ckpt_dir / "ve.pt", weights_only=True)
+        # )
+        ve.load_state_dict(torch.load(
+        ckpt_dir / "ve.pt", 
+        weights_only=True, 
+        map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    ))
         ve.to(device).eval()
 
         t3 = T3(T3Config.multilingual())
@@ -175,9 +180,14 @@ class ChatterboxMultilingualTTS:
         t3.to(device).eval()
 
         s3gen = S3Gen()
-        s3gen.load_state_dict(
-            torch.load(ckpt_dir / "s3gen.pt", weights_only=True)
-        )
+        # s3gen.load_state_dict(
+        #     torch.load(ckpt_dir / "s3gen.pt", weights_only=True)
+        # )
+        s3gen.load_state_dict(torch.load(
+            ckpt_dir / "s3gen.pt", 
+            weights_only=True, 
+            map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        ))
         s3gen.to(device).eval()
 
         tokenizer = MTLTokenizer(