EvoTransformer-v2.1 / init_save.py
HemanM's picture
Create init_save.py
03f2ff2 verified
raw
history blame
442 Bytes
# init_save.py
import os
import torch
from evo_model import EvoTransformerForClassification, EvoTransformerConfig
# βœ… Ensure the directory exists
os.makedirs("trained_model", exist_ok=True)
# Create config and model instance
config = EvoTransformerConfig()
model = EvoTransformerForClassification(config)
# Save initial model and config
model.save_pretrained("trained_model")
print("βœ… Initial EvoTransformer saved to trained_model/")