"""List selected H3 checkpoint tensor shapes."""

from safetensors import safe_open


path = "/models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
with safe_open(path, framework="pt", device="cpu") as checkpoint:
    for name in checkpoint.keys():
        if name == "adaln_t_table" or name.startswith("blocks.0.adaln_proj") or name.startswith("final_layer.adaln_proj"):
            print(name, tuple(checkpoint.get_tensor(name).shape))