Fixed flatten_parameters() warning for models/tacotron.py

Author: Ryan Butler

models/tacotron.py

@@ -206,17 +206,30 @@ def forward(self, encoder_seq_proj, query, t):
 class Decoder(nn.Module):
     def __init__(self, n_mels, decoder_dims, lstm_dims):
         super().__init__()
+
+        # List of all rnns to call `flatten_parameters()` on
+        self._to_flatten = []
+
         self.max_r = 20
         self.r = None
         self.generating = False
         self.n_mels = n_mels
         self.prenet = PreNet(n_mels)
         self.attn_net = LSA(decoder_dims)
+        
         self.attn_rnn = nn.GRUCell(decoder_dims + decoder_dims // 2, decoder_dims)
+        self._to_flatten.append(self.attn_rnn)
+
         self.rnn_input = nn.Linear(2 * decoder_dims, lstm_dims)
+
         self.res_rnn1 = nn.LSTMCell(lstm_dims, lstm_dims)
         self.res_rnn2 = nn.LSTMCell(lstm_dims, lstm_dims)
+        self._to_flatten += [self.res_rnn1, self.res_rnn2]
+
         self.mel_proj = nn.Linear(lstm_dims, n_mels * self.max_r, bias=False)
+
+        # Avoid fragmentation of RNN parameters and associated warning
+        self._flatten_parameters()
 
     def zoneout(self, prev, current, p=0.1):
         device = next(self.parameters()).device  # Use same device as parameters
@@ -225,6 +238,11 @@ def zoneout(self, prev, current, p=0.1):
 
     def forward(self, encoder_seq, encoder_seq_proj, prenet_in, 
                 hidden_states, cell_states, context_vec, t):
+
+        # Although we `_flatten_parameters()` on init, when using DataParallel 
+        # the model gets replicated, making it no longer guaranteed that the
+        # weights are contiguous in GPU memory. Hence, we must call it again
+        self._flatten_parameters()
 
         # Need this for reshaping mels
         batch_size = encoder_seq.size(0)