Assign module to device after quantization

[jcaip]

Summary:

Before, we were moving the module to the device and then quantizing it, now we quantize first and move to the device after.

This was causing some of the huggingface integration tests to fail.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm failures during provisioning step due to network issues

❌ 2 New Failures

As of commit 31e7381 with merge base 2a89491 ():

NEW FAILURES - The following jobs have failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  test_to_device_granularity1_sizes2

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