import torch
import io
model = torch.nn.Sequential(torch.nn.Linear(32, 256, dtype=torch.bfloat16, device="cuda"))
from torchao.quantization import Int4WeightOnlyConfig, quantize_
quant_config = Int4WeightOnlyConfig(group_size=128, packing_format="plain", version=2)
quantize_(model, quant_config)
buf = io.BytesIO()
torch.save(model.state_dict(), buf)
# Push to hub
USER_ID = "torchao-testing"
MODEL_NAME = "single-linear"
save_to = f"{USER_ID}/{MODEL_NAME}-INT4-v2-0.13-dev"
print("save to:", save_to)
# Manual Testing
from huggingface_hub import HfApi
api = HfApi()
api.create_repo(save_to, repo_type="model", exist_ok=True)
api.upload_file(
path_or_fileobj=buf,
path_in_repo="model.bin",
repo_id=save_to,
)