import copy

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(7)

features = torch.linspace(-1.5, 1.5, steps=24, dtype=torch.float32).reshape(8, 3)
targets = features @ torch.tensor([[0.7], [-0.2], [0.4]]) + 0.15

microbatch_size = 2
accumulation_steps = 2
effective_batch_size = microbatch_size * accumulation_steps

train_loader = DataLoader(
    TensorDataset(features, targets),
    batch_size=microbatch_size,
    shuffle=False,
)

model = nn.Linear(3, 1)
reference_model = copy.deepcopy(model)
initial_parameters = [parameter.detach().clone() for parameter in model.parameters()]

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.05)