Как использовать набор данных numpy в Pytorch Lightning

Я хочу создать набор данных, используя NumPy, а затем хочу обучить и протестировать простую модель, такую ​​как «линейная или логистическая».

Я пытаюсь выучить Pytorch Lightning. Я нашел руководство, в котором мы можем использовать Набор данных NumPy и может использовать единое распределение здесь. Как новичок, я не понимаю, как я могу это сделать!

Мой код приведен ниже

import numpy as np 
import pytorch_lightning as pl 
from torch.utils.data import random_split, DataLoader, TensorDataset
import torch
from torch.autograd import Variable
from torchvision import transforms

np.random.seed(42)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

class DataModuleClass(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
        self.constant = 2
        self.batch_size = 10
        self.transform = transforms.Compose([
            transforms.ToTensor()
        ])
        
    def prepare_data(self):
        a = np.random.uniform(0, 500, 500)
        b = np.random.normal(0, self.constant, len(x))

        c = a + b
        X = np.transpose(np.array([a, b]))

        idx = np.arange(500)
        np.random.shuffle(idx)
        
        # Uses foirst 400 random indices for training
        train_idx = idx[:400]
        # Uses the remaining indices for validation
        val_idx = idx[400:]
        
        # Generate train and validation dataset
        x_train, y_train = X[train_idx], y[train_idx]
        x_val, y_val = X[val_idx], y[val_idx]
        
        # Converting numpy array to Tensor
        self.x_train_tensor = torch.from_numpy(x_train).float().to(device)
        self.y_train_tensor = torch.from_numpy(y_train).float().to(device)
        
        self.x_val_tensor = torch.from_numpy(x_val).float().to(device)
        self.y_val_tensor = torch.from_numpy(y_val).float().to(device)
        
        training_dataset = TensorDataset(self.x_train_tensor, self.y_train_tensor)
        
        validation_dataset = TensorDataset(self.x_val_tensor, self.y_val_tensor)

        return training_dataset, validation_dataset
        
    def train_dataloader(self):
        training_dataloader = prepare_data() # Most probably this is wrong way!!!
        return DataLoader(self.training_dataloader)

    def val_dataloader(self):
        validation_dataloader = prepare_data() # Most probably this is wrong way!!!
        return DataLoader(self.validation_dataloader)
        
    # def test_dataloader(self):
        
obj = DataModuleClass()
print(obj.prepare_data())  

Эта часть выполняется на основе предоставленного ответа [Здесь я хочу принять a and b as features и c как label or target variable.]

Теперь, как я могу передать набор данных в метод обучения и проверки?