pass self.device to load_fsspec function

In case using cuda by default, pass self.device to choose the right device user defined in to() function.
coqui-ai · Sep 18, 2023 · eff8817 · eff8817
1 parent f829bf5
commit eff8817
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/TTS/tts/models/xtts.py b/TTS/tts/models/xtts.py
@@ -642,7 +642,7 @@ def load_checkpoint(
         self.init_models()
         if eval:
             self.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache)
-        self.load_state_dict(load_fsspec(model_path)["model"], strict=strict)
+        self.load_state_dict(load_fsspec(model_path, map_location=self.device)["model"], strict=strict)
 
         if eval:
             self.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache)