Fix failing tests on h100

[jainapurva]

Updating tests and code to fix current issues in running tests on H100.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 5 Cancelled Jobs

As of commit 2832ec5 with merge base d72a6d1 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 728379b0cc8a6b0e5e10a62251b848f816c98bf998bce09a397498329b73ddc6 /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 8ce6f09186bd1097f443cca2a2a6a030c03b9c070fb750e960af6e31a2aa19aa /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t b8c2579e3b8b7acb560ea4319771a7bdc291b679c53af91497b9e78d4478cbf5 /exec failed with exit code 1

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• Code Analysis with Ruff / build (3.9) (gh)
• PR Label Check / Check PR Labels (gh)
• Run Float8 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio --index-url https://download.pytor... / linux-job (gh)
  ##[error]The operation was canceled.
• Run Float8 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  ##[error]The operation was canceled.
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)

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