modify unflatten for vllm

[liangel-02]

Summary

We keep track of what has been processed by deleting processed keys from the original state dict. This is for the case when not all tensor attributes are available to us when loading, so we reconstruct the tensor subclasses with missing attributes in a later call.

We expect the state dict to be empty after loading is complete

Test
python test/prototype/safetensors/test_safetensors_support.py

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3297

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit cdc5ca2 with merge base 6c78c4d ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.