DeepTrust / convert_model.py
priyansh-nagar's picture
Upload 5 files
80cf891 verified
raw
history blame contribute delete
418 Bytes
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")