tiny-recursive-model 0.0.4__tar.gz → 0.0.6__tar.gz

{tiny_recursive_model-0.0.4 → tiny_recursive_model-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tiny-recursive-model
-Version: 0.0.4
+Version: 0.0.6
 Summary: Tiny Recursive Model
 Project-URL: Homepage, https://pypi.org/project/tiny-recursive-model/
 Project-URL: Repository, https://github.com/lucidrains/tiny-recursive-model
@@ -38,7 +38,7 @@ Requires-Dist: accelerate
 Requires-Dist: einops>=0.8.1
 Requires-Dist: ema-pytorch
 Requires-Dist: torch>=2.4
-Requires-Dist: x-transformers
+Requires-Dist: x-transformers>=2.8.4
 Provides-Extra: examples
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
@@ -47,7 +47,7 @@ Description-Content-Type: text/markdown
 
 <img width="300" alt="trm-fig1" src="https://github.com/user-attachments/assets/950db79e-5f9c-4fec-a4e4-7b9355b39ce8" />
 
-## Tiny Recursive Model (TRM) wip
+## Tiny Recursive Model (TRM)
 
 Implementation of [Tiny Recursive Model](https://arxiv.org/abs/2510.04871) (TRM), improvement to [HRM](https://github.com/lucidrains/hrm) from Sapient AI, by [Alexia Jolicoeur-Martineau](https://ajolicoeur.wordpress.com/about/)
 

{tiny_recursive_model-0.0.4 → tiny_recursive_model-0.0.6}/README.md

@@ -1,7 +1,7 @@
 
 <img width="300" alt="trm-fig1" src="https://github.com/user-attachments/assets/950db79e-5f9c-4fec-a4e4-7b9355b39ce8" />
 
-## Tiny Recursive Model (TRM) wip
+## Tiny Recursive Model (TRM)
 
 Implementation of [Tiny Recursive Model](https://arxiv.org/abs/2510.04871) (TRM), improvement to [HRM](https://github.com/lucidrains/hrm) from Sapient AI, by [Alexia Jolicoeur-Martineau](https://ajolicoeur.wordpress.com/about/)
 

{tiny_recursive_model-0.0.4 → tiny_recursive_model-0.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tiny-recursive-model"
-version = "0.0.4"
+version = "0.0.6"
 description = "Tiny Recursive Model"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,7 +27,7 @@ dependencies = [
     "einops>=0.8.1",
     "ema-pytorch",
     "torch>=2.4",
-    "x-transformers",
+    "x-transformers>=2.8.4",
 ]
 
 [project.urls]

{tiny_recursive_model-0.0.4 → tiny_recursive_model-0.0.6}/tiny_recursive_model/trainer.py

@@ -86,7 +86,7 @@ class Trainer(Module):
 
         self.optim = optim_klass(
             model.parameters(),
-            lr = learning_rate,
+            lr = learning_rate / (batch_size * max_recurrent_steps),
             weight_decay = weight_decay
         )
 

{tiny_recursive_model-0.0.4 → tiny_recursive_model-0.0.6}/tiny_recursive_model/trm.py

@@ -220,4 +220,4 @@ class TinyRecursiveModel(Module):
 
         losses = (loss, halt_loss)
 
-        return (total_loss.mean(), losses, *return_package)
+        return (total_loss.sum(), losses, *return_package)