mirror of https://github.com/InternLM/InternLM
8 lines
259 B
Python
8 lines
259 B
Python
|
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
|