The batch dimension is usually in dim0 and the DataLoader will use the default_collate method to create a batch of samples using this structure.
Based on your code snippet it looks like the batch dimension is in dim1, which yields this behavior.
If you don’t want to flatten your data, you could permute the dimensions:
class ModelDataset(Dataset):
    
    def __init__(self, data, labels, ind):

        # take only the data specified by ind. ind is a list of indexes
        self.data = data[:,ind,:].permute(1, 0, 2)  # size = (ind, 6, 256)
        self.labels = labels[ind]
       
    def __getitem__(self, index):
        
        pair = self.data[index] # pair = (6,256)
        label = self.labels[index]
        
        x1 = pair[:, :128]
        x2 = pair[:, 128:]
        
        return x1.T, x2.T, label
    
    def __len__(self):
        return len(self.data)

Let me know, if that works for you!