|
from transformers import AutoModel
|
|
|
|
model = AutoModel.from_pretrained("../hf_ckpt/", trust_remote_code=True).cuda()
|
|
print(model)
|
|
assert model.config.hidden_size == 2048
|
|
assert model.config.num_attention_heads == 16
|
|
assert model.config.num_hidden_layers == 16
|