File size: 508 Bytes
27c649c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | # utils/data_loader.py
import torch
from torch.utils.data import Dataset, DataLoader
class CustomDataset(Dataset):
def __init__(self, data):
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx]
def load_data(batch_size=32):
# Dummy data
data = [torch.randn(10) for _ in range(1000)]
dataset = CustomDataset(data)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return loader
|