def load_mlp_128_256(num_features, **_):
    """Two layer MLP backbone with 128 and 256 hidden units respectively."""
    return (
        nn.Sequential(
            nn.Linear(num_features, 128),
            nn.ReLU(True),
            nn.BatchNorm1d(128),
            nn.Linear(128, 256),
            nn.ReLU(True),
            nn.BatchNorm1d(256),
        ),
        256,
    )
Two layer MLP backbone with 128 and 256 hidden units respectively.
src/flash/template/classification/backbones.py