allow for turning off norm in random projection quantizer

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.1.4',
+  version = '1.1.5',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/random_projection_quantizer.py

@@ -15,6 +15,7 @@ def __init__(
         codebook_size,
         codebook_dim,
         num_codebooks = 1,
+        norm = True,
         **kwargs
     ):
         super().__init__()
@@ -28,7 +29,7 @@ def __init__(
         # in section 3 of https://arxiv.org/abs/2202.01855
         # "The input data is normalized to have 0 mean and standard deviation of 1 ... to prevent collapse"
 
-        self.norm = nn.LayerNorm(dim, elementwise_affine = False)
+        self.norm = nn.LayerNorm(dim, elementwise_affine = False) if norm else nn.Identity()
 
         self.vq = VectorQuantize(
             dim = codebook_dim * num_codebooks,