text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  class Flatten(nn.Module):
  
      def __init__(self):
          super(Flatten, self).__init__()
  
      def forward(self, x):
          return x.view(x.shape[0], -1)
  
  
  class Concatenate(nn.Module):
  
      def __init__(self, dim=-1):
          super(Concatenate, self).__init__()
          self.dim = dim
  
      def forward(self, x):
          return torch.cat(x, dim=self.dim)
  