How to use pytorch DataLoader with a 3-D matrix for LSTM input?
I have a dataset of 3-D(time_stepinputsizetotal_num) matrix which is a .mat file. I want to use DataLoader to get a input dataset for LSTM which batch_size is 5. My code is as following:

file_path = "…/database/frameLength100/notOverlap/a.mat"
mat_data = s.loadmat(file_path)
tensor_data = torch.from_numpy(mat_data[‘a’]) #Tensor

class CustomDataset(Dataset):

def __init__(self, tensor_data):
    self.tensor_data = tensor_data

def __getitem__(self, index):
    data = self.tensor_data[index]
    label = 1;
    return data, label

def __len__(self):
    return len(self.tensor_data)
custom_dataset = CustomDataset(tensor_data=tensor_data)
train_loader = DataLoader(dataset=custom_dataset, batch_size=5, shuffle=True)


I think the code is wrong but I have no idea how to correct it. What makes me confused is how how can I make DataLoader know which dimension is ‘total_num’ so that I get the dataset which batch size is 5.

If I understand correctly, you want the batching to happen along the total_num dimension, i. e. dimension 2.

You could simply use that the dimension to index your dataset, i.e. change __getitem__ to data = self.tensor_data[:, :, index], and accordingly in __len__, return self.tensor_data.size(2) instead of len(self.tensor_data). Each batch will then have size [time_step, inputsize, 5].
