Tacotron model JIT works?

tacotron_melgan/tacotron_model.py

@@ -255,8 +255,8 @@ class Decoder(nn.Module):
         decoder_input: all zeros frames
         """
         B = memory.size(0)
-        decoder_input = Variable(memory.data.new(
-            B, self.n_mel_channels * self.n_frames_per_step).zero_())
+        mem = memory.data.new(B, self.n_mel_channels * self.n_frames_per_step)
+        decoder_input = Variable(mem.zero_())
         return decoder_input
 
     def initialize_decoder_states(self, memory, mask):
@@ -548,4 +548,5 @@ class Tacotron2(nn.Module):
               alignments
             ])
 
-        return outputs
+        #return outputs
+        return outputs[0]

tacotron_melgan/wrap_tacotron2_model.py

@@ -76,15 +76,18 @@ module = Tacotron2(HParams())
 
 print('Load state dict...')
 module.load_state_dict(tacotron_model['state_dict'])
-#module.eval(inference=False)
+module.eval() # NB: Complains unless called `Did you forget call .eval() on your model?`
 
-print('JIT model...')
-mel_file = '/home/bt/dev/voder/data/mels/LJ002-0320.mel'
-example = torch.load(mel_file, map_location=torch.device('cpu')) # TODO
-traced_script_module = torch.jit.trace(module, example)
+#print('JIT model...')
+text_sequence_file = '/home/bt/dev/voder/data/text/tacotron_text_sequence.pt'
+#text_sequence_file = '/home/bt/dev/voder/data/mels/LJ002-0320.mel'
+text_sequence = torch.load(text_sequence_file, map_location=torch.device('cpu'))
+# NB: Getting `Tracing failed sanity checks! Graphs differed across invocations!`
+traced_script_module = torch.jit.trace(module, text_sequence, check_trace=False)
 traced_script_module.save("tacotron_container.pt")
 
 #print('Saving model...')
 #container = torch.jit.script(module)
 #container.save("tacotron_container.pt")
 
+print('Done.')