class_docstr: Given an input shaped like ``(batch_size, time_steps, [rest])`` and a ``Module`` that takes
  inputs like ``(batch_size, [rest])``, ``TimeDistributed`` reshapes the input to be
  ``(batch_size * time_steps, [rest])``, applies the contained ``Module``, then reshapes it back.
  
  Note that while the above gives shapes with ``batch_size`` first, this ``Module`` also works if
  ``batch_size`` is second - we always just combine the first two dimensions, then split them.
class_docstr_tok: ["Given","an","input","shaped","like","``","(","batch_size",",","time_steps",",","[","rest","]",")","``","and","a","``","Module","``","that","takes","inputs","like","``","(","batch_size",",","[","rest","]",")","``",",","``","TimeDistributed","``","reshapes","the","input","to","be","``","(","batch_size","*","time_steps",",","[","rest","]",")","``",",","applies","the","contained","``","Module","``",",","then","reshapes","it","back",".","Note","that","while","the","above","gives","shapes","with","``","batch_size","``","first",",","this","``","Module","``","also","works","if","``","batch_size","``","is","second","-","we","always","just","combine","the","first","two","dimensions",",","then","split","them","."]
code: def forward(self, *inputs):  # pylint: disable=arguments-differ
      reshaped_inputs = []
      for input_tensor in inputs:
          input_size = input_tensor.size()
          if len(input_size) <= 2:
              raise RuntimeError(
                  "No dimension to distribute: " + str(input_size)
              )
  
          # Squash batch_size and time_steps into a single axis; result has shape
          # (batch_size * time_steps, input_size).
          squashed_shape = [-1] + [x for x in input_size[2:]]
          reshaped_inputs.append(
              input_tensor.contiguous().view(*squashed_shape)
          )
  
      reshaped_outputs = self._module(*reshaped_inputs)
  
      # Now get the output back into the right shape.
      # (batch_size, time_steps, [hidden_size])
      # noinspection PyUnboundLocalVariable
      new_shape = [input_size[0], input_size[1]] + [
          x for x in reshaped_outputs.size()[1:]
      ]
      outputs = reshaped_outputs.contiguous().view(*new_shape)
  
      return outputs
  
file_path: jiant/jiant/ext/allennlp.py
imports: from typing import Optional
  import torch
  import torch.nn
label_desc: Use image node_id 2 for calling the TimeDistributed obj's underlying member method code with example usage: obj.forward() and returns: outputs