vector-quantize-pytorch 1.17.0__tar.gz → 1.17.2__tar.gz

{vector_quantize_pytorch-1.17.0 → vector_quantize_pytorch-1.17.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.17.0
+Version: 1.17.2
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch

{vector_quantize_pytorch-1.17.0 → vector_quantize_pytorch-1.17.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.17.0"
+version = "1.17.2"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{vector_quantize_pytorch-1.17.0 → vector_quantize_pytorch-1.17.2}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -187,7 +187,7 @@ class FSQ(Module):
         # whether to force quantization step to be full precision or not
 
         force_f32 = self.force_quantization_f32
-        quantization_context = partial(autocast, enabled = False) if force_f32 else nullcontext
+        quantization_context = partial(autocast, 'cuda', enabled = False) if force_f32 else nullcontext
 
         with quantization_context():
             orig_dtype = z.dtype

{vector_quantize_pytorch-1.17.0 → vector_quantize_pytorch-1.17.2}/vector_quantize_pytorch/lookup_free_quantization.py

@@ -10,7 +10,9 @@ from math import log2, ceil
 from functools import partial, cache
 from collections import namedtuple
 from contextlib import nullcontext
+
 import torch.distributed as dist
+from torch.distributed import nn as dist_nn
 
 import torch
 from torch import nn, einsum
@@ -36,7 +38,7 @@ def maybe_distributed_mean(t):
     if not is_distributed():
         return t
 
-    dist.nn.all_reduce(t)
+    dist_nn.all_reduce(t)
     t = t / dist.get_world_size()
     return t