diff --git a/vall_e/train.py b/vall_e/train.py
index a292635..5260e3b 100644
--- a/vall_e/train.py
+++ b/vall_e/train.py
@@ -83,7 +83,7 @@ def main():
                 resps_list = model(
                     text_list=batch["text"],
                     proms_list=batch["proms"],
-                    resp_list=batch["resp"],
+                    resps_list=[r.unsqueeze(-1) for r in batch["resp"]],
                     sampling_temperature=cfg.sampling_temperature,
                 )
             else: