# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""AlexNet."""

from torch.nn import BCEWithLogitsLoss
from torch.optim import SGD

from ptbench.data.augmentations import ElasticDeformation
from ptbench.models.alexnet import Alexnet

model = Alexnet(
    train_loss=BCEWithLogitsLoss(),
    validation_loss=BCEWithLogitsLoss(),
    optimizer_type=SGD,
    optimizer_arguments=dict(lr=0.01, momentum=0.1),
    augmentation_transforms=[ElasticDeformation(p=0.8)],
    pretrained=False,
)