address #145

lucidrains · Jul 1, 2024 · 1bce1c3 · 1bce1c3
1 parent 1889d04
commit 1bce1c3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.14.44"
+version = "1.14.45"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

diff --git a/vector_quantize_pytorch/lookup_free_quantization.py b/vector_quantize_pytorch/lookup_free_quantization.py
@@ -270,7 +270,8 @@ def forward(
 
         assert x.shape[-1] == self.dim, f'expected dimension of {self.dim} but received {x.shape[-1]}'
 
-        x = self.project_in(x)
+        with autocast():
+            x = self.project_in(x)
 
         # maybe soft clamp
 
@@ -382,7 +383,8 @@ def forward(
 
         # project out to feature dimension if needed
 
-        x = self.project_out(x)
+        with autocast():
+            x = self.project_out(x)
 
         # reconstitute image or video dimensions