Update test_quantized_training.py

msaroufim · web-flow · commit ec23da937354 · 2025-05-01T12:34:44.000-07:00
diff --git a/test/prototype/test_quantized_training.py b/test/prototype/test_quantized_training.py
@@ -30,7 +30,7 @@
     Transformer,
 )
 
-from torchao.low_bit_optim import _AdamW
+from torchao.optim import _AdamW
 from torchao.prototype.quantized_training import (
     Int8MixedPrecisionTrainingConfig,
     bitnet_training,

Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@`
`30`	`30`	`Transformer,`
`31`	`31`	`)`
`32`	`32`
`33`		`-from torchao.low_bit_optim import _AdamW`
	`33`	`+from torchao.optim import _AdamW`
`34`	`34`	`from torchao.prototype.quantized_training import (`
`35`	`35`	`Int8MixedPrecisionTrainingConfig,`
`36`	`36`	`bitnet_training,`