How to use the babi.BABIDataset function in babi

def __init__(self, X, q, y):
        assert X.shape[0] == q.shape[0]
        assert X.shape[0] == y.shape[0]
        
        self.X = X
        self.q = q
        self.y = y
    
    def __getitem__(self, idx):
        return (self.X[idx], self.q[idx]), self.y[idx]
    
    def __len__(self):
        return self.X.shape[0]

train_data = BABIDataset(X=X_train, q=q_train, y=y_train)
test_data = BABIDataset(X=X_test, q=q_test, y=y_test)

train_indices, search_indices = train_test_split(range(len(X_train)), train_size=0.5)
dataloaders = {
    "train"  : ZipDataloader([
        torch.utils.data.DataLoader(
            dataset=train_data,
            batch_size=32,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices),
        ),
        torch.utils.data.DataLoader(
            dataset=train_data,
            batch_size=32,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(search_indices),
        )
    ]),
    "test"  : DataLoader(

class BABIDataset(Dataset):
    def __init__(self, X, q, y):
        assert X.shape[0] == q.shape[0]
        assert X.shape[0] == y.shape[0]
        
        self.X = X
        self.q = q
        self.y = y
    
    def __getitem__(self, idx):
        return (self.X[idx], self.q[idx]), self.y[idx]
    
    def __len__(self):
        return self.X.shape[0]

train_data = BABIDataset(X=X_train, q=q_train, y=y_train)
test_data = BABIDataset(X=X_test, q=q_test, y=y_test)

train_indices, search_indices = train_test_split(range(len(X_train)), train_size=0.5)
dataloaders = {
    "train"  : ZipDataloader([
        torch.utils.data.DataLoader(
            dataset=train_data,
            batch_size=32,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices),
        ),
        torch.utils.data.DataLoader(
            dataset=train_data,
            batch_size=32,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(search_indices),
        )
    ]),