"""TensorViz public example: a residual classifier with a 128-feature block."""
import torch.nn as nn


class ResidualBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.expand = nn.Linear(128, 256)
        self.relu = nn.ReLU()
        self.project = nn.Linear(256, 128)

    def forward(self, x):
        residual = x
        x = self.expand(x)
        x = self.relu(x)
        x = self.project(x)
        return x + residual


# tensorviz: input=1,64
class ResidualClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Linear(64, 128)
        self.block = ResidualBlock()
        self.head = nn.Linear(64, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.block(x)
        return self.head(x)