Tacotron2 --cpu-run fix

shutovilya · shutovilya · commit 680ccfcf89a9 · 2020-05-21T21:05:24.000+07:00
diff --git a/PyTorch/SpeechSynthesis/Tacotron2/inference.py b/PyTorch/SpeechSynthesis/Tacotron2/inference.py
@@ -148,15 +148,15 @@ def pad_sequences(batch):
     return text_padded, input_lengths
 
 
-def prepare_input_sequence(texts):
+def prepare_input_sequence(texts, cpu_run=False):
 
     d = []
     for i,text in enumerate(texts):
         d.append(torch.IntTensor(
             text_to_sequence(text, ['english_cleaners'])[:]))
 
     text_padded, input_lengths = pad_sequences(d)
-    if torch.cuda.is_available():
+    if torch.cuda.is_available() and not cpu_run:
         text_padded = torch.autograd.Variable(text_padded).cuda().long()
         input_lengths = torch.autograd.Variable(input_lengths).cuda().long()
     else:
@@ -236,7 +236,7 @@ def main():
 
     measurements = {}
 
-    sequences_padded, input_lengths = prepare_input_sequence(texts)
+    sequences_padded, input_lengths = prepare_input_sequence(texts, args.cpu_run)
 
     with torch.no_grad(), MeasureTime(measurements, "tacotron2_time", args.cpu_run):
         mel, mel_lengths, alignments = jitted_tacotron2(sequences_padded, input_lengths)
diff --git a/PyTorch/SpeechSynthesis/Tacotron2/test_infer.py b/PyTorch/SpeechSynthesis/Tacotron2/test_infer.py
@@ -193,7 +193,7 @@ def main():
         measurements = {}
 
         with MeasureTime(measurements, "pre_processing", args.cpu_run):
-            sequences_padded, input_lengths = prepare_input_sequence(texts)
+            sequences_padded, input_lengths = prepare_input_sequence(texts, args.cpu_run)
 
         with torch.no_grad():
             with MeasureTime(measurements, "latency", args.cpu_run):
