pranavKHF/code_vulnerability_detection

1import torch2from src.model import VulnerabilityCodeT53 4# Load original big checkpoint5checkpoint = torch.load("models/best_model.pt", map_location="cpu")6 7# Initialize model8model = VulnerabilityCodeT5(num_labels=2)9 10# Load only model weights11model.load_state_dict(checkpoint['model_state_dict'])12 13# Save clean weights only14torch.save(model.state_dict(), "models/best_model_clean.pt")15 16print("Saved clean model.")17