dream-s1k-demo / list_linear.py
况兑
feat: add Dream loss probe + linear module lister
85a4f1f
raw
history blame contribute delete
372 Bytes
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)