neulab/docprompting-conala huggingface dataset:
torch.nn.functional.pad(input, pad, mode='constant', value=0)  
Pads tensor.  Padding size:

The padding size by which to pad some dimensions of input are described starting from the last dimension and moving forward. ⌊len(pad)2⌋\left\lfloor\frac{\text{len(pad)}}{2}\right\rfloor  dimensions of input will be padded. For example, to pad only the last dimension of the input tensor, then pad has the form (padding_left,padding_right)(\text{padding\_left}, \text{padding\_right}) ; to pad the last 2 dimensions of the input tensor, then use (padding_left,padding_right,(\text{padding\_left}, \text{padding\_right},  padding_top,padding_bottom)\text{padding\_top}, \text{padding\_bottom}) ; to pad the last 3 dimensions, use (padding_left,padding_right,(\text{padding\_left}, \text{padding\_right},  padding_top,padding_bottom\text{padding\_top}, \text{padding\_bottom}  padding_front,padding_back)\text{padding\_front}, \text{padding\_back}) .  Padding mode:

See torch.nn.ConstantPad2d, torch.nn.ReflectionPad2d, and torch.nn.ReplicationPad2d for concrete examples on how each of the padding modes works. Constant padding is implemented for arbitrary dimensions. Replicate padding is implemented for padding the last 3 dimensions of 5D input tensor, or the last 2 dimensions of 4D input tensor, or the last dimension of 3D input tensor. Reflect padding is only implemented for padding the last 2 dimensions of 4D input tensor, or the last dimension of 3D input tensor.    Note When using the CUDA backend, this operation may induce nondeterministic behaviour in its backward pass that is not easily switched off. Please see the notes on Reproducibility for background.   Parameters 
 
input (Tensor) – N-dimensional tensor 
pad (tuple) – m-elements tuple, where m2≤\frac{m}{2} \leq  input dimensions and mm  is even. 
mode – 'constant', 'reflect', 'replicate' or 'circular'. Default: 'constant'
 
value – fill value for 'constant' padding. Default: 0
    Examples: >>> t4d = torch.empty(3, 3, 4, 2)
>>> p1d = (1, 1) # pad last dim by 1 on each side
>>> out = F.pad(t4d, p1d, "constant", 0)  # effectively zero padding
>>> print(out.size())
torch.Size([3, 3, 4, 4])
>>> p2d = (1, 1, 2, 2) # pad last dim by (1, 1) and 2nd to last by (2, 2)
>>> out = F.pad(t4d, p2d, "constant", 0)
>>> print(out.size())
torch.Size([3, 3, 8, 4])
>>> t4d = torch.empty(3, 3, 4, 2)
>>> p3d = (0, 1, 2, 1, 3, 3) # pad by (0, 1), (2, 1), and (3, 3)
>>> out = F.pad(t4d, p3d, "constant", 0)
>>> print(out.size())
torch.Size([3, 9, 7, 3])
torch.nn.functional.pairwise_distance(x1, x2, p=2.0, eps=1e-06, keepdim=False) [source]
 
See torch.nn.PairwiseDistance for details
torch.nn.functional.pdist(input, p=2) → Tensor  
Computes the p-norm distance between every pair of row vectors in the input. This is identical to the upper triangular portion, excluding the diagonal, of torch.norm(input[:, None] - input, dim=2, p=p). This function will be faster if the rows are contiguous. If input has shape N×MN \times M  then the output will have shape 12N(N−1)\frac{1}{2} N (N - 1) . This function is equivalent to scipy.spatial.distance.pdist(input, ‘minkowski’, p=p) if p∈(0,∞)p \in (0, \infty) . When p=0p = 0  it is equivalent to scipy.spatial.distance.pdist(input, ‘hamming’) * M. When p=∞p = \infty , the closest scipy function is scipy.spatial.distance.pdist(xn, lambda x, y: np.abs(x - y).max()).  Parameters 
 
input – input tensor of shape N×MN \times M . 
p – p value for the p-norm distance to calculate between each vector pair ∈[0,∞]\in [0, \infty] .
torch.nn.functional.pixel_shuffle(input, upscale_factor) → Tensor  
Rearranges elements in a tensor of shape (∗,C×r2,H,W)(*, C \times r^2, H, W)  to a tensor of shape (∗,C,H×r,W×r)(*, C, H \times r, W \times r) , where r is the upscale_factor. See PixelShuffle for details.  Parameters 
 
input (Tensor) – the input tensor 
upscale_factor (int) – factor to increase spatial resolution by    Examples: >>> input = torch.randn(1, 9, 4, 4)
>>> output = torch.nn.functional.pixel_shuffle(input, 3)
>>> print(output.size())
torch.Size([1, 1, 12, 12])
torch.nn.functional.pixel_unshuffle(input, downscale_factor) → Tensor  
Reverses the PixelShuffle operation by rearranging elements in a tensor of shape (∗,C,H×r,W×r)(*, C, H \times r, W \times r)  to a tensor of shape (∗,C×r2,H,W)(*, C \times r^2, H, W) , where r is the downscale_factor. See PixelUnshuffle for details.  Parameters 
 
input (Tensor) – the input tensor 
downscale_factor (int) – factor to increase spatial resolution by    Examples: >>> input = torch.randn(1, 1, 12, 12)
>>> output = torch.nn.functional.pixel_unshuffle(input, 3)
>>> print(output.size())
torch.Size([1, 9, 4, 4])
torch.nn.functional.poisson_nll_loss(input, target, log_input=True, full=False, size_average=None, eps=1e-08, reduce=None, reduction='mean') [source]
 
Poisson negative log likelihood loss. See PoissonNLLLoss for details.  Parameters 
 
input – expectation of underlying Poisson distribution. 
target – random sample target∼Poisson(input)target \sim \text{Poisson}(input) . 
log_input – if True the loss is computed as exp⁡(input)−target∗input\exp(\text{input}) - \text{target} * \text{input} , if False then loss is input−target∗log⁡(input+eps)\text{input} - \text{target} * \log(\text{input}+\text{eps}) . Default: True
 
full – whether to compute full loss, i. e. to add the Stirling approximation term. Default: False target∗log⁡(target)−target+0.5∗log⁡(2∗π∗target)\text{target} * \log(\text{target}) - \text{target} + 0.5 * \log(2 * \pi * \text{target}) . 
size_average (bool, optional) – Deprecated (see reduction). By default, the losses are averaged over each loss element in the batch. Note that for some losses, there multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True
 
eps (float, optional) – Small value to avoid evaluation of log⁡(0)\log(0)  when log_input`=``False`. Default: 1e-8 
reduce (bool, optional) – Deprecated (see reduction). By default, the losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True
 
reduction (string, optional) – Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'
torch.nn.functional.prelu(input, weight) → Tensor [source]
 
Applies element-wise the function PReLU(x)=max⁡(0,x)+weight∗min⁡(0,x)\text{PReLU}(x) = \max(0,x) + \text{weight} * \min(0,x)  where weight is a learnable parameter. See PReLU for more details.
torch.nn.functional.relu(input, inplace=False) → Tensor [source]
 
Applies the rectified linear unit function element-wise. See ReLU for more details.
torch.nn.functional.relu6(input, inplace=False) → Tensor [source]
 
Applies the element-wise function ReLU6(x)=min⁡(max⁡(0,x),6)\text{ReLU6}(x) = \min(\max(0,x), 6) . See ReLU6 for more details.
torch.nn.functional.relu_(input) → Tensor  
In-place version of relu().
torch.nn.functional.rrelu(input, lower=1./8, upper=1./3, training=False, inplace=False) → Tensor [source]
 
Randomized leaky ReLU. See RReLU for more details.
torch.nn.functional.rrelu_(input, lower=1./8, upper=1./3, training=False) → Tensor  
In-place version of rrelu().
torch.nn.functional.selu(input, inplace=False) → Tensor [source]
 
Applies element-wise, SELU(x)=scale∗(max⁡(0,x)+min⁡(0,α∗(exp⁡(x)−1)))\text{SELU}(x) = scale * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1))) , with α=1.6732632423543772848170429916717\alpha=1.6732632423543772848170429916717  and scale=1.0507009873554804934193349852946scale=1.0507009873554804934193349852946 . See SELU for more details.
torch.nn.functional.sigmoid(input) → Tensor [source]
 
Applies the element-wise function Sigmoid(x)=11+exp⁡(−x)\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}  See Sigmoid for more details.
torch.nn.functional.silu(input, inplace=False) [source]
 
Applies the silu function, element-wise.  silu(x)=x∗σ(x),where σ(x) is the logistic sigmoid.\text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}  
 Note See Gaussian Error Linear Units (GELUs) where the SiLU (Sigmoid Linear Unit) was originally coined, and see Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning and Swish: a Self-Gated Activation Function where the SiLU was experimented with later.  See SiLU for more details.
torch.nn.functional.smooth_l1_loss(input, target, size_average=None, reduce=None, reduction='mean', beta=1.0) [source]
 
Function that uses a squared term if the absolute element-wise error falls below beta and an L1 term otherwise. See SmoothL1Loss for details.
torch.nn.functional.softmax(input, dim=None, _stacklevel=3, dtype=None) [source]
 
Applies a softmax function. Softmax is defined as: Softmax(xi)=exp⁡(xi)∑jexp⁡(xj)\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}  It is applied to all slices along dim, and will re-scale them so that the elements lie in the range [0, 1] and sum to 1. See Softmax for more details.  Parameters 
 
input (Tensor) – input 
dim (int) – A dimension along which softmax will be computed. 
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.     Note This function doesn’t work directly with NLLLoss, which expects the Log to be computed between the Softmax and itself. Use log_softmax instead (it’s faster and has better numerical properties).
torch.nn.functional.softmin(input, dim=None, _stacklevel=3, dtype=None) [source]
 
Applies a softmin function. Note that Softmin(x)=Softmax(−x)\text{Softmin}(x) = \text{Softmax}(-x) . See softmax definition for mathematical formula. See Softmin for more details.  Parameters 
 
input (Tensor) – input 
dim (int) – A dimension along which softmin will be computed (so every slice along dim will sum to 1). 
dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.
torch.nn.functional.softplus(input, beta=1, threshold=20) → Tensor  
Applies element-wise, the function Softplus(x)=1β∗log⁡(1+exp⁡(β∗x))\text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x)) . For numerical stability the implementation reverts to the linear function when input×β>thresholdinput \times \beta > threshold . See Softplus for more details.
torch.nn.functional.softshrink(input, lambd=0.5) → Tensor  
Applies the soft shrinkage function elementwise See Softshrink for more details.
torch.nn.functional.softsign(input) → Tensor [source]
 
Applies element-wise, the function SoftSign(x)=x1+∣x∣\text{SoftSign}(x) = \frac{x}{1 + |x|}  See Softsign for more details.
torch.nn.functional.soft_margin_loss(input, target, size_average=None, reduce=None, reduction='mean') → Tensor [source]
 
See SoftMarginLoss for details.
torch.nn.functional.tanh(input) → Tensor [source]
 
Applies element-wise, Tanh(x)=tanh⁡(x)=exp⁡(x)−exp⁡(−x)exp⁡(x)+exp⁡(−x)\text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}  See Tanh for more details.
torch.nn.functional.tanhshrink(input) → Tensor [source]
 
Applies element-wise, Tanhshrink(x)=x−Tanh(x)\text{Tanhshrink}(x) = x - \text{Tanh}(x)  See Tanhshrink for more details.
torch.nn.functional.threshold(input, threshold, value, inplace=False)  
Thresholds each element of the input Tensor. See Threshold for more details.
torch.nn.functional.threshold_(input, threshold, value) → Tensor  
In-place version of threshold().
torch.nn.functional.triplet_margin_loss(anchor, positive, negative, margin=1.0, p=2, eps=1e-06, swap=False, size_average=None, reduce=None, reduction='mean') [source]
 
See TripletMarginLoss for details
torch.nn.functional.triplet_margin_with_distance_loss(anchor, positive, negative, *, distance_function=None, margin=1.0, swap=False, reduction='mean') [source]
 
See TripletMarginWithDistanceLoss for details.
torch.nn.functional.unfold(input, kernel_size, dilation=1, padding=0, stride=1) [source]
 
Extracts sliding local blocks from a batched input tensor.  Warning Currently, only 4-D input tensors (batched image-like tensors) are supported.   Warning More than one element of the unfolded tensor may refer to a single memory location. As a result, in-place operations (especially ones that are vectorized) may result in incorrect behavior. If you need to write to the tensor, please clone it first.  See torch.nn.Unfold for details
torch.nn.functional.upsample(input, size=None, scale_factor=None, mode='nearest', align_corners=None) [source]
 
Upsamples the input to either the given size or the given scale_factor  Warning This function is deprecated in favor of torch.nn.functional.interpolate(). This is equivalent with nn.functional.interpolate(...).   Note This operation may produce nondeterministic gradients when given tensors on a CUDA device. See Reproducibility for more information.  The algorithm used for upsampling is determined by mode. Currently temporal, spatial and volumetric upsampling are supported, i.e. expected inputs are 3-D, 4-D or 5-D in shape. The input dimensions are interpreted in the form: mini-batch x channels x [optional depth] x [optional height] x width. The modes available for upsampling are: nearest, linear (3D-only), bilinear, bicubic (4D-only), trilinear (5D-only)  Parameters 
 
input (Tensor) – the input tensor 
size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]) – output spatial size. 
scale_factor (float or Tuple[float]) – multiplier for spatial size. Has to match input size if it is a tuple. 
mode (string) – algorithm used for upsampling: 'nearest' | 'linear' | 'bilinear' | 'bicubic' | 'trilinear'. Default: 'nearest'
 
align_corners (bool, optional) – Geometrically, we consider the pixels of the input and output as squares rather than points. If set to True, the input and output tensors are aligned by the center points of their corner pixels, preserving the values at the corner pixels. If set to False, the input and output tensors are aligned by the corner points of their corner pixels, and the interpolation uses edge value padding for out-of-boundary values, making this operation independent of input size when scale_factor is kept the same. This only has an effect when mode is 'linear', 'bilinear', 'bicubic' or 'trilinear'. Default: False
     Note With mode='bicubic', it’s possible to cause overshoot, in other words it can produce negative values or values greater than 255 for images. Explicitly call result.clamp(min=0, max=255) if you want to reduce the overshoot when displaying the image.   Warning With align_corners = True, the linearly interpolating modes (linear, bilinear, and trilinear) don’t proportionally align the output and input pixels, and thus the output values can depend on the input size. This was the default behavior for these modes up to version 0.3.1. Since then, the default behavior is align_corners = False. See Upsample for concrete examples on how this affects the outputs.
torch.nn.functional.upsample_bilinear(input, size=None, scale_factor=None) [source]
 
Upsamples the input, using bilinear upsampling.  Warning This function is deprecated in favor of torch.nn.functional.interpolate(). This is equivalent with nn.functional.interpolate(..., mode='bilinear', align_corners=True).  Expected inputs are spatial (4 dimensional). Use upsample_trilinear fo volumetric (5 dimensional) inputs.  Parameters 
 
input (Tensor) – input 
size (int or Tuple[int, int]) – output spatial size. 
scale_factor (int or Tuple[int, int]) – multiplier for spatial size     Note This operation may produce nondeterministic gradients when given tensors on a CUDA device. See Reproducibility for more information.
torch.nn.functional.upsample_nearest(input, size=None, scale_factor=None) [source]
 
Upsamples the input, using nearest neighbours’ pixel values.  Warning This function is deprecated in favor of torch.nn.functional.interpolate(). This is equivalent with nn.functional.interpolate(..., mode='nearest').  Currently spatial and volumetric upsampling are supported (i.e. expected inputs are 4 or 5 dimensional).  Parameters 
 
input (Tensor) – input 
size (int or Tuple[int, int] or Tuple[int, int, int]) – output spatia size. 
scale_factor (int) – multiplier for spatial size. Has to be an integer.     Note This operation may produce nondeterministic gradients when given tensors on a CUDA device. See Reproducibility for more information.
class torch.nn.GaussianNLLLoss(*, full=False, eps=1e-06, reduction='mean') [source]
 
Gaussian negative log likelihood loss. The targets are treated as samples from Gaussian distributions with expectations and variances predicted by the neural network. For a D-dimensional target tensor modelled as having heteroscedastic Gaussian distributions with a D-dimensional tensor of expectations input and a D-dimensional tensor of positive variances var the loss is:  loss=12∑i=1D(log⁡(max(var[i], eps))+(input[i]−target[i])2max(var[i], eps))+const.\text{loss} = \frac{1}{2}\sum_{i=1}^D \left(\log\left(\text{max}\left(\text{var}[i], \ \text{eps}\right)\right) + \frac{\left(\text{input}[i] - \text{target}[i]\right)^2} {\text{max}\left(\text{var}[i], \ \text{eps}\right)}\right) + \text{const.}  
where eps is used for stability. By default, the constant term of the loss function is omitted unless full is True. If var is a scalar (implying target tensor has homoscedastic Gaussian distributions) it is broadcasted to be the same size as the input.  Parameters 
 
full (bool, optional) – include the constant term in the loss calculation. Default: False. 
eps (float, optional) – value used to clamp var (see note below), for stability. Default: 1e-6. 
reduction (string, optional) – specifies the reduction to apply to the output:'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the output is the average of all batch member losses, 'sum': the output is the sum of all batch member losses. Default: 'mean'.     Shape:

 Input: (N,∗)(N, *)  where ∗*  means any number of additional dimensions Target: (N,∗)(N, *) , same shape as the input Var: (N,1)(N, 1)  or (N,∗)(N, *) , same shape as the input Output: scalar if reduction is 'mean' (default) or 'sum'. If reduction is 'none', then (N)(N) 
    Examples: >>> loss = nn.GaussianNLLLoss()
>>> input = torch.randn(5, 2, requires_grad=True)
>>> target = torch.randn(5, 2)
>>> var = torch.ones(5, 2, requires_grad=True) #heteroscedastic
>>> output = loss(input, target, var)
>>> output.backward()


>>> loss = nn.GaussianNLLLoss()
>>> input = torch.randn(5, 2, requires_grad=True)
>>> target = torch.randn(5, 2)
>>> var = torch.ones(5, 1, requires_grad=True) #homoscedastic
>>> output = loss(input, target, var)
>>> output.backward()
  Note The clamping of var is ignored with respect to autograd, and so the gradients are unaffected by it.   Reference:

Nix, D. A. and Weigend, A. S., “Estimating the mean and variance of the target probability distribution”, Proceedings of 1994 IEEE International Conference on Neural Networks (ICNN’94), Orlando, FL, USA, 1994, pp. 55-60 vol.1, doi: 10.1109/ICNN.1994.374138.
class torch.nn.GELU [source]
 
Applies the Gaussian Error Linear Units function:  GELU(x)=x∗Φ(x)\text{GELU}(x) = x * \Phi(x)  
where Φ(x)\Phi(x)  is the Cumulative Distribution Function for Gaussian Distribution.  Shape:

 Input: (N,∗)(N, *)  where * means, any number of additional dimensions Output: (N,∗)(N, *) , same shape as the input     Examples: >>> m = nn.GELU()
>>> input = torch.randn(2)
>>> output = m(input)
class torch.nn.GroupNorm(num_groups, num_channels, eps=1e-05, affine=True) [source]
 
Applies Group Normalization over a mini-batch of inputs as described in the paper Group Normalization  y=x−E[x]Var[x]+ϵ∗γ+βy = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta  
The input channels are separated into num_groups groups, each containing num_channels / num_groups channels. The mean and standard-deviation are calculated separately over the each group. γ\gamma  and β\beta  are learnable per-channel affine transform parameter vectors of size num_channels if affine is True. The standard-deviation is calculated via the biased estimator, equivalent to torch.var(input, unbiased=False). This layer uses statistics computed from input data in both training and evaluation modes.  Parameters 
 
num_groups (int) – number of groups to separate the channels into 
num_channels (int) – number of channels expected in input 
eps – a value added to the denominator for numerical stability. Default: 1e-5 
affine – a boolean value that when set to True, this module has learnable per-channel affine parameters initialized to ones (for weights) and zeros (for biases). Default: True.     Shape:

 Input: (N,C,∗)(N, C, *)  where C=num_channelsC=\text{num\_channels} 
 Output: (N,C,∗)(N, C, *)  (same shape as input)    Examples: >>> input = torch.randn(20, 6, 10, 10)
>>> # Separate 6 channels into 3 groups
>>> m = nn.GroupNorm(3, 6)
>>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm)
>>> m = nn.GroupNorm(6, 6)
>>> # Put all 6 channels into a single group (equivalent with LayerNorm)
>>> m = nn.GroupNorm(1, 6)
>>> # Activating the module
>>> output = m(input)
class torch.nn.GRUCell(input_size, hidden_size, bias=True) [source]
 
A gated recurrent unit (GRU) cell  r=σ(Wirx+bir+Whrh+bhr)z=σ(Wizx+biz+Whzh+bhz)n=tanh⁡(Winx+bin+r∗(Whnh+bhn))h′=(1−z)∗n+z∗h\begin{array}{ll} r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\ z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\ n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\ h' = (1 - z) * n + z * h \end{array} 
where σ\sigma  is the sigmoid function, and ∗*  is the Hadamard product.  Parameters 
 
input_size – The number of expected features in the input x
 
hidden_size – The number of features in the hidden state h
 
bias – If False, then the layer does not use bias weights b_ih and b_hh. Default: True
     Inputs: input, hidden

 
input of shape (batch, input_size): tensor containing input features 
hidden of shape (batch, hidden_size): tensor containing the initial hidden state for each element in the batch. Defaults to zero if not provided.   Outputs: h’

 
h’ of shape (batch, hidden_size): tensor containing the next hidden state for each element in the batch   Shape:

 Input1: (N,Hin)(N, H_{in})  tensor containing input features where HinH_{in}  = input_size
 Input2: (N,Hout)(N, H_{out})  tensor containing the initial hidden state for each element in the batch where HoutH_{out}  = hidden_size Defaults to zero if not provided. Output: (N,Hout)(N, H_{out})  tensor containing the next hidden state for each element in the batch     Variables 
 
~GRUCell.weight_ih – the learnable input-hidden weights, of shape (3*hidden_size, input_size)
 
~GRUCell.weight_hh – the learnable hidden-hidden weights, of shape (3*hidden_size, hidden_size)
 
~GRUCell.bias_ih – the learnable input-hidden bias, of shape (3*hidden_size)
 
~GRUCell.bias_hh – the learnable hidden-hidden bias, of shape (3*hidden_size)
     Note All the weights and biases are initialized from U(−k,k)\mathcal{U}(-\sqrt{k}, \sqrt{k})  where k=1hidden_sizek = \frac{1}{\text{hidden\_size}}   Examples: >>> rnn = nn.GRUCell(10, 20)
>>> input = torch.randn(6, 3, 10)
>>> hx = torch.randn(3, 20)
>>> output = []
>>> for i in range(6):
        hx = rnn(input[i], hx)
        output.append(hx)