text: import torch
  
  
  class Flattener(torch.nn.Module):
      def __init__(self):
          """
          Flattens last 3 dimensions to make it only batch size, -1
          """
          super(Flattener, self).__init__()
  
      def forward(self, x):
          return x.view(x.size(0), -1)
  