Hi All,
I am trying to implement custom LSTM layer with custom cell.
It is working OK when I pass only one sample, but when I want to pass a batch of data a problem appear.
For example:
If my data look like [128, 64] where 128 is max sequence length and 64 is word embedding size, the model works fine
I want to know how to handle a batch of data e.g. [16, 128, 64] ?
I search for the implementation of the lstm cell in pytorch , but I can’t find the code that handles the batch
  

      github.com
  

  
    pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/torch/nn/modules/rnn.py#L689



      
      
                                    max_batch_size, self.hidden_size,
                                    dtype=input.dtype, device=input.device)
              hx = (h_zeros, c_zeros)
          else:
              # Each batch of the hidden state should match the input sequence that
              # the user believes he/she is passing in.
              hx = self.permute_hidden(hx, sorted_indices)
          
          
self.check_forward_args(input, hx, batch_sizes)
          if batch_sizes is None:
              result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
                                self.dropout, self.training, self.bidirectional, self.batch_first)
          else:
              result = _VF.lstm(input, batch_sizes, hx, self._flat_weights, self.bias,
                                self.num_layers, self.dropout, self.training, self.bidirectional)
          output = result[0]
          hidden = result[1:]
          # xxx: isinstance check needs to be in conditional for TorchScript to compile
          if isinstance(orig_input, PackedSequence):
              output_packed = PackedSequence(output, batch_sizes, sorted_indices, unsorted_indices)
              return output_packed, self.permute_hidden(hidden, unsorted_indices)
      
    
  


  

  
    
    
  

  


and this code leads to this impl.
  

      github.com
  

  
    pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/test/test_jit.py#L2640



      
      
                  return x + 10
          
          
    with TemporaryFileName() as fname:
                  fn.save(fname)
                  with io.open(fname, 'rb') as f:
                      self.assertTrue(torch.serialization._is_zipfile(f))
          
          
def test_python_bindings(self):
              lstm_cell = torch.jit.script(LSTMCellS)
          
          
    def lstm(x, hx, cx, w_ih, w_hh, b_ih, b_hh):
                  for i in range(x.size(0)):
                      hx, cx = lstm_cell(x[i], hx, cx, w_ih, w_hh, b_ih, b_hh)
                  return hx
          
          
    slstm = torch.jit.script(lstm)
          
          
    inputs = get_lstm_inputs('cpu', training=True, seq_length=10)
              slstm(*inputs).sum().backward()
              global fw_graph
              fw_graph = slstm.graph_for(*inputs)
      
    
  


  

  
    
    
  

  


is pytorch for loop the data batch?
Thanks in advance.