import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.models import ResNet18_Weights, resnet18


torch.manual_seed(29)

weights = ResNet18_Weights.DEFAULT
train_data = ImageFolder(
    "data/document-images/train",
    transform=weights.transforms(),
)
train_loader = DataLoader(train_data, batch_size=4, shuffle=True)
class_names = train_data.classes