able to override reasoning steps on forward

Author: lucidrains

HRM/hrm.py

@@ -3,7 +3,7 @@
 
 import torch
 import torch.nn.functional as F
-from torch import nn, Tensor, tensor, is_tensor, cat, stack
+from torch import Tensor, tensor, is_tensor, cat, stack
 from torch.nn import Embedding, Linear, Module, ModuleList
 from torch.utils._pytree import tree_map
 
@@ -38,9 +38,9 @@ def __init__(
         super().__init__()
         self.num_hiddens_to_concat = num_hiddens_to_concat
 
-        self.norms = ModuleList([nn.RMSNorm(dim) for _ in range(num_hiddens_to_concat)])
+        self.norms = ModuleList([RMSNorm(dim) for _ in range(num_hiddens_to_concat)])
 
-        self.to_combined = nn.Linear(dim * self.num_hiddens_to_concat, dim, bias = False)
+        self.to_combined = Linear(dim * self.num_hiddens_to_concat, dim, bias = False)
 
     def forward(
         self,
@@ -132,9 +132,12 @@ def forward(
         *,
         labels = None,
         detach_hiddens = True,
-        one_step_grad = True
+        one_step_grad = True,
+        reasoning_steps = None
     ):
 
+        reasoning_steps = default(reasoning_steps, self.reasoning_steps)
+
         if detach_hiddens:
             hiddens = tree_map_tensor(hiddens, lambda t: t.detach())
 
@@ -184,7 +187,7 @@ def evaluate_network_(
         context = torch.no_grad if one_step_grad else nullcontext
 
         with context():
-            for index in range(self.reasoning_steps * self.lowest_steps_per_reasoning_step - 1):
+            for index in range(reasoning_steps * self.lowest_steps_per_reasoning_step - 1):
                 iteration = index + 1
 
                 for network_index, (network, hidden_combine, evaluate_network_at) in enumerate(zip(self.networks, self.hidden_combiners, self.evaluate_networks_at)):

tests/test_hrm.py

@@ -30,6 +30,7 @@ def test_hrm():
         ],
         num_tokens = 256,
         dim = 512,
+        reasoning_steps = 3
 
     )
 
@@ -38,3 +39,7 @@ def test_hrm():
 
     loss, (logits, hiddens) = hrm(seq, labels = labels)
     loss.backward()
+
+    # after much training
+
+    pred = hrm(seq, reasoning_steps = 5)