import torch
from torchvision import models

model = models.efficientnet_b0(pretrained=False)
model.classifier[1] = torch.nn.Linear(
    model.classifier[1].in_features, 2
)

state_dict = torch.load(
    "models/efficientnet_b0_ffpp_c23.pth",
    map_location="cpu"
)
model.load_state_dict(state_dict, strict=False)

torch.save(model, "models/deeptrust.pt")
print("Model saved as models/deeptrust.pt")