# Load model directly
from transformers import AutoModel
model = AutoModel.from_pretrained("Zyphra/BlackMamba-2.8B")