<reponame>busycalibrating/Adversarial-Training
import torch.nn as nn
from adv_train.utils import Flatten


def build_model_mnist(device=None):
    model = nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1),
        nn.ReLU(),
        nn.Conv2d(32, 32, 3, padding=1, stride=2),
        nn.ReLU(),
        nn.Conv2d(32, 64, 3, padding=1),
        nn.ReLU(),
        nn.Conv2d(64, 64, 3, padding=1, stride=2),
        nn.ReLU(),
        Flatten(),
        nn.Linear(7 * 7 * 64, 100),
        nn.ReLU(),
        nn.Linear(100, 10),
    )
    model = model.to(device)
    return model

