x-transformers 2.11.16__tar.gz → 2.11.17__tar.gz

{x_transformers-2.11.16 → x_transformers-2.11.17}/PKG-INFO

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

{x_transformers-2.11.16 → x_transformers-2.11.17}/pyproject.toml

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

{x_transformers-2.11.16 → x_transformers-2.11.17}/x_transformers/free_transformer.py

@@ -149,6 +149,7 @@ class FreeTransformer(Module):
         enc_kwargs: dict = dict(),
         dec_kwargs: dict = dict(),
         kl_loss_weight = 1.,
+        latent_dropout_prob = 0.,
         pad_id = -1,
         **kwargs
     ):
@@ -187,6 +188,8 @@ class FreeTransformer(Module):
 
         self.from_latent_to_condition = nn.Linear(self.binary_mapper.num_codes, dim, bias = False)
 
+        self.latent_dropout = nn.Dropout(latent_dropout_prob)
+
         self.decoder_head = Decoder(
             dim = dim,
             depth = dec_head_depth,
@@ -380,6 +383,8 @@ class FreeTransformer(Module):
 
         latents, kl_loss = self.encode_to_latents(tokens_for_latents, mask = encoder_mask, per_token_latents = per_token_latents, return_kl_loss = True)
 
+        latents = self.latent_dropout(latents)
+
         condition = self.from_latent_to_condition(latents)
 
         # decoder tail