import torch
from torch import nn, optim

class M(nn.Module):
    def __init__(self, depth: int, dim: int, hidden_dim: Optional[int] = None):
        super().__init__()
        hidden_dim = hidden_dim or dim
        self.layers = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(dim, hidden_dim),
                    nn.BatchNorm1d(hidden_dim),
                    nn.LeakyReLU(),
                )
            ]
            + [
                nn.Sequential(
                    nn.Linear(hidden_dim, hidden_dim),
                    nn.BatchNorm1d(hidden_dim),
                    nn.LeakyReLU(),
                )
                for i in range(depth)
            ]
            + [nn.Linear(hidden_dim, dim), nn.Sigmoid()]
        )

    def forward(self, batch: torch.Tensor) -> torch.Tensor:
        x = batch
        for l in self.layers:
            x = l(x)

        return x


model = M(depth=5, dim=40).train() # on CPU

from layer_to_layer_pytorch.l2l import Layer2Layer

l2l_model = Layer2Layer(
    model,
    layers_attr="layers", # attribute with ModuleList
    microbatch_size=100,  # size of a microbatch in a minibatch :) from original paper
    verbose=False  # enable tqdm
)

from tqdm.auto import tqdm, trange

x = torch.rand(1_000, 40) # on CPU
y = torch.rand(1_000, 40) # on CPU

losses = []
criterion = nn.MSELoss()

optimizer = optim.AdamW(l2l_model.main_params) # optimizer works with the main model on CPU

for i in trange(2000):
    l2l_model.zero_grad()
    _ = l2l_model.forward(x)

    loss_value: float = l2l_model.compute_loss(y, criterion)

    if i % 50 == 0:
        tqdm.write(f"[{i}] loss = {loss_value}")
    losses.append(loss_value)


    l2l_model.backward()
    optimizer.step()
    l2l_model.update_main_model_params() # Sync params with CPU

from layer_to_layer_pytorch.l2l import Layer2Layer

l2l_model = Layer2Layer(
    model,
    layers_attr="layers",
    microbatch_size=100,

    # fp-16
    mixed_precision=True,
    loss_scale = 128.0
)