class_docstr: Wrapper for torch.nn.LSTM that handles masked inputs.
class_docstr_tok: ["Wrapper","for","torch.nn.LSTM","that","handles","masked","inputs","."]
docstr: Args:
      data (torch.Tensor): tensor with shape (batch_size, seq_len, feature_dim)
      mask (torch.Tensor): bool tensor with shape (batch_size, seq_len).
          Sequence elements that are False are invalid.
  
  Returns:
      Tensor with shape (batch_size, output_dim) - outoput_dim is determined by
          hidden_dim and whether bidirectional or not
docstr_tok: ["Args",":","data","(","torch.Tensor",")",":","tensor","with","shape","(","batch_size",",","seq_len",",","feature_dim",")","mask","(","torch.Tensor",")",":","bool","tensor","with","shape","(","batch_size",",","seq_len",")",".","Sequence","elements","that","are","False","are","invalid",".","Returns",":","Tensor","with","shape","(","batch_size",",","output_dim",")","-","outoput_dim","is","determined","by","hidden_dim","and","whether","bidirectional","or","not"]
code: def forward(
      self, data: torch.Tensor, mask: Optional[torch.Tensor] = None
  ) -> torch.Tensor:
      """
      Args:
          data (torch.Tensor): tensor with shape (batch_size, seq_len, feature_dim)
          mask (torch.Tensor): bool tensor with shape (batch_size, seq_len).
              Sequence elements that are False are invalid.
  
      Returns:
          Tensor with shape (batch_size, output_dim) - outoput_dim is determined by
              hidden_dim and whether bidirectional or not
      """
      assert data.dim() == 3
      b, t = data.shape[0], data.shape[1]
  
      if mask is None:
          mask = torch.ones((b, t), dtype=torch.bool)
  
      lengths = mask.sum(axis=1)
      x_packed = pack_padded_sequence(
          data,
          lengths.clamp(1, data.size(1)),
          batch_first=True,
          enforce_sorted=False,
      )
      _, (h, _) = self.lstm(x_packed)
  
      if self.bidirectional:
          out = torch.cat([h[0, :, :], h[1, :, :]], dim=-1)
      else:
          out = h[-1, :, :]
  
      return out
  
imports: from typing import List, Optional, Tuple
  import torch
  from pytorchvideo.layers.utils import set_attributes
  from torch import nn
  from torch.nn.utils.rnn import pack_padded_sequence
label_desc: Use image node_id 2 for calling the LSTM obj's underlying member method code with example usage: obj.forward(data, mask) and returns: out