import torch


def reshape_data(data_list, batch_first=True, collapse_dims=False):
    """Concat input/output data and clear buffers

    Args:
      data_list (list): list of tensors of equal, arbitrary shape where the
          batch_dim is either 0 or 1 depending on self.batch_first.
      batch_first (bool, optional): is batch dim first. (default: True)
      collapse_dim (bool, optional): if True, collapse all but the last dim
          together forming a 2D output tensor.

    Returns:
      Single tensor with all tensors from data_list concatenated across
      batch_dim. Guarenteed to be 2D if collapse_dims=True.
    """
    d = torch.cat(data_list, dim=int(not batch_first))
    if collapse_dims and len(d.shape) > 2:
        d = d.view(-1, d.shape[-1])
    return d