Merge pull request #2189 from coqui-ai/fix-capacitron-test

2022-12-06 21:25:06 +01:00 · 2022-12-06 21:25:06 +01:00 · c753ad49cc
parent 9321b22203 4787a2a993
commit c753ad49cc
1 changed files with 1 additions and 1 deletions
--- a/tests/tts_tests/test_tacotron2_model.py
+++ b/tests/tts_tests/test_tacotron2_model.py
@ -301,7 +301,7 @@ class TacotronCapacitronTrainTest(unittest.TestCase):
        batch["stop_targets"] = (batch["stop_targets"].sum(2) > 0.0).unsqueeze(2).float().squeeze()

        model = Tacotron2(config).to(device)
-        criterion = model.get_criterion()
+        criterion = model.get_criterion().to(device)
        optimizer = model.get_optimizer()

        model.train()