File size: 372 Bytes
85a4f1f
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
import os, torch
from transformers import AutoModel
MODEL_ID = os.getenv("MODEL_ID", "Dream-org/Dream-v0-Instruct-7B")
REV = os.getenv("REV", None)
m = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, revision=REV)
print("[INFO] Listing torch.nn.Linear modules:")
for name, mod in m.named_modules():
    if isinstance(mod, torch.nn.Linear):
        print(name)