source_current/synthetic_task/data.py

import torch
from models import MLP

def genData(device, input_dim=64, output_dim=10, num_samples=2048, batch_size=32):
    # Generate some data
    y = torch.rand(num_samples, output_dim).to(device)
    model = MLP(input_dim=output_dim, output_dim=input_dim).to(device)

    x = model(y).detach() 

    # Add some noise
    x = x + torch.rand(num_samples, input_dim).to(device) * 0.1

    dataset = list(zip(x,y))
    train_split = 0.8
    training_size = int(num_samples * train_split)
    
    train_loader = torch.utils.data.DataLoader(dataset[:training_size], batch_size=batch_size, shuffle=True)
    test_loader  = torch.utils.data.DataLoader(dataset[training_size:], batch_size=batch_size, shuffle=False)
    return train_loader, test_loader