code: def merge_first_two_dims(tensor):
      
      shape = tensor.shape
      batch_size = shape[0] * shape[1]
      new_shape = tuple([batch_size] + list(shape[2:]))
      return tensor.view(new_shape)
docs: Reshape tensor to merge first two dimensions