x-transformers 2.4.6__tar.gz → 2.4.7__tar.gz

{x_transformers-2.4.6 → x_transformers-2.4.7}/PKG-INFO

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

{x_transformers-2.4.6 → x_transformers-2.4.7}/pyproject.toml

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

{x_transformers-2.4.6 → x_transformers-2.4.7}/tests/test_x_transformers.py

@@ -1100,7 +1100,10 @@ def add_attn_pool():
 
     assert intermediates.attn_pooled_tokens.shape[1] == 3
 
-def test_up():
+@pytest.mark.parametrize('keep_buffer_on_cpu', (False, True))
+def test_up(
+    keep_buffer_on_cpu
+):
     from x_transformers.up_wrapper import UniversalPretrainWrapper
 
     model = TransformerWrapper(
@@ -1115,7 +1118,11 @@ def test_up():
         ),
     )
 
-    up_wrapper = UniversalPretrainWrapper(model, seq_len = 16)
+    up_wrapper = UniversalPretrainWrapper(
+        model,
+        seq_len = 16,
+        keep_buffer_on_cpu = keep_buffer_on_cpu
+    )
 
     loss = up_wrapper()
     loss.backward()

{x_transformers-2.4.6 → x_transformers-2.4.7}/x_transformers/up_wrapper.py

@@ -146,7 +146,8 @@ class UniversalPretrainWrapper(Module):
         batch_size = 32,
         seq_len = 512,
         seed_length = 8,
-        reset_turing_machine_every = 0
+        reset_turing_machine_every = 0,
+        keep_buffer_on_cpu = False
     ):
         super().__init__()
 
@@ -185,12 +186,16 @@ class UniversalPretrainWrapper(Module):
 
         init_data_buffer = self.random_sequences_fn(buffer_size // 2, buffer_size // 2)
 
-        self.register_buffer('synth_data_buffer', init_data_buffer)
+        if keep_buffer_on_cpu:
+            self.synth_data_buffer = init_data_buffer
+        else:
+            self.register_buffer('synth_data_buffer', init_data_buffer)
+
         self.register_buffer('step', tensor(0))
 
     @property
     def device(self):
-        return self.synth_data_buffer.device
+        return self.step.device
 
     def get_rand_sequences_from_buffer(self, size = None):
         size = default(size, self.batch_size)
@@ -217,8 +222,8 @@ class UniversalPretrainWrapper(Module):
 
         generated = self.data_generator.generate(
             self.seq_len,
-            condition = conditions,
-            seed = seeds
+            condition = conditions.to(self.device),
+            seed = seeds.to(self.device)
         )
 
         self.step.add_(1)
@@ -244,6 +249,6 @@ class UniversalPretrainWrapper(Module):
 
         # sample yet again according to pseudocode
 
-        data = self.get_rand_sequences_from_buffer()
+        data = self.get_rand_sequences_from_buffer().to(self.device)
 
         return self.ar_wrapped(data)