EvoTransformer-v2.1 / init_model.py
HemanM's picture
Update init_model.py
3570d2d verified
raw
history blame
431 Bytes
import os
import torch
from evo_model import EvoTransformerConfig, EvoTransformerForClassification
# Ensure folder
os.makedirs("trained_model", exist_ok=True)
# Create config and model
config = EvoTransformerConfig()
model = EvoTransformerForClassification(config)
# Save config and model in Hugging Face format
model.save_pretrained("trained_model")
print("✅ EvoTransformer initial model + config saved to 'trained_model/'")