Fix for weights-only load (pytorch#1228)

stack-info: PR: pytorch#1228, branch: drisspg/stack/19
nighting0le01 · Dec 5, 2024 · 9ec0edd · 9ec0edd
1 parent 87a1e65
commit 9ec0edd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
@@ -368,8 +368,9 @@ def test_optim_bf16_stochastic_round_correctness(self):
             optim2.step()
             optim2.zero_grad()
 
-            torch.testing.assert_close(loss1, loss2, msg=lambda msg: f"Iteration {idx}. {msg}")
-
+            torch.testing.assert_close(
+                loss1, loss2, msg=lambda msg: f"Iteration {idx}. {msg}"
+            )
 
 
 _FSDP_WORLD_SIZE = 2

diff --git a/torchao/prototype/low_bit_optim/adam.py b/torchao/prototype/low_bit_optim/adam.py
@@ -13,12 +13,12 @@
 
 class _AdamBase(Optimizer):
     def __init__(
-        self,
-        params,
-        lr,
-        betas,
-        eps,
-        weight_decay,
+        self, 
+        params, 
+        lr, 
+        betas, 
+        eps, 
+        weight_decay, 
         amsgrad,
         *,
         block_size,