From 8d7f9685a554a236cbed524a5cf5444091692491 Mon Sep 17 00:00:00 2001
From: Driss Guessous <32754868+drisspg@users.noreply.github.com>
Date: Tue, 5 Nov 2024 19:10:28 -0800
Subject: [PATCH] Fix for weights-only load (#1228)

stack-info: PR: https://github.com/pytorch/ao/pull/1228, branch: drisspg/stack/19
---
 test/prototype/test_low_bit_optim.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
index 71518047a1..c87f482dbb 100644
--- a/test/prototype/test_low_bit_optim.py
+++ b/test/prototype/test_low_bit_optim.py
@@ -19,8 +19,11 @@
     quantize_4bit_with_qmap,
     _fp32_to_bf16_sr,
 )
-
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_3, TORCH_VERSION_AT_LEAST_2_4, TORCH_VERSION_AT_LEAST_2_6
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_3,
+    TORCH_VERSION_AT_LEAST_2_4,
+    TORCH_VERSION_AT_LEAST_2_6,
+)
 
 try:
     import bitsandbytes as bnb