matching parameters counts with constructor of optimStateFp8

pytorch · Nov 12, 2024 · 38951ae · 38951ae
1 parent 366743c
commit 38951ae
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchao/prototype/low_bit_optim/subclass_fp8.py b/torchao/prototype/low_bit_optim/subclass_fp8.py
@@ -42,7 +42,7 @@ class OptimStateFp8(TorchAOBaseTensor):
     tensor_attrs = ["codes", "scale"]
 
     @staticmethod
-    def __new__(cls, codes: Tensor, scale: Tensor):
+    def __new__(cls, codes: Tensor, scale: Tensor, k: Optional[Tensor] =None):
         return Tensor._make_wrapper_subclass(cls, codes.shape, device=codes.device)
 
     def __init__(self, codes: Tensor, scale: Tensor, k: Optional[Tensor] =None):