Merge pull request #139 from BenjaminBossan/fix-type-annotation-torch…

…-2.5 FIX: Type annotation torch.Any
Vahe1994 · Feb 26, 2025 · 429b86e · 429b86e
2 parents ab272bf + 780fe15
commit 429b86e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/inference_lib/src/aqlm/inference.py b/inference_lib/src/aqlm/inference.py
@@ -1,6 +1,6 @@
 """ Core mathematics for Additive Quantization (AQ): initialization, reconstruction and beam search"""
 import math
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 import torch.nn as nn
@@ -99,7 +99,7 @@ def _get_autograd_matmul_op(forward_pass_kernel, backward_pass_kernel):
     class _QuantizedMatmul(torch.autograd.Function):
         @staticmethod
         def forward(
-            ctx: torch.Any,
+            ctx: Any,
             input: torch.Tensor,
             codes: torch.IntTensor,
             codebooks: torch.Tensor,