x-transformers 2.8.1__tar.gz → 2.8.2__tar.gz

{x_transformers-2.8.1 → x_transformers-2.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.8.1
+Version: 2.8.2
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.8.1 → x_transformers-2.8.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.8.1"
+version = "2.8.2"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.8.1 → x_transformers-2.8.2}/x_transformers/gpt_vae.py

@@ -46,25 +46,29 @@ class GPTVAE(Module):
         vae_kl_loss_weight = 1.,
         latents_dropout_prob = 0.5, # what percentage of the time to dropout the latents completely
         pad_id = -1,
+        encoder: Module | None = None,
         **kwargs
     ):
         super().__init__()
         dim_latent = default(dim_latent, dim)
 
-        self.encoder = TransformerWrapper(
-            num_tokens = num_tokens,
-            max_seq_len = max_seq_len + 1,
-            return_only_embed = True,
-            average_pool_embed = True,
-            attn_layers = Encoder(
-                dim = dim,
-                depth = enc_depth,
-                attn_dim_head = attn_dim_head,
-                heads = heads,
-                **kwargs,
-                **enc_kwargs
-            ),
-        )
+        if not exists(encoder):
+            encoder = TransformerWrapper(
+                num_tokens = num_tokens,
+                max_seq_len = max_seq_len + 1,
+                return_only_embed = True,
+                average_pool_embed = True,
+                attn_layers = Encoder(
+                    dim = dim,
+                    depth = enc_depth,
+                    attn_dim_head = attn_dim_head,
+                    heads = heads,
+                    **kwargs,
+                    **enc_kwargs
+                ),
+            )
+
+        self.encoder = encoder
 
         self.to_latent_mean_log_variance = nn.Sequential(
             nn.Linear(dim, dim_latent * 2),