another step

Author: lucidrains

HRM/hrm.py

@@ -71,6 +71,7 @@ def __init__(
     ):
         super().__init__()
         attn_layers_klass = Encoder if not causal else Decoder
+        self.causal = causal
 
         # input
 
@@ -133,9 +134,14 @@ def forward(
         labels = None,
         detach_hiddens = True,
         one_step_grad = True,
-        reasoning_steps = None
+        reasoning_steps = None,
+        return_autoreg_loss = False
     ):
 
+        if return_autoreg_loss:
+            assert self.causal
+            seq, labels = seq[:, :-1], seq[:, 1:]
+
         return_loss = exists(labels)
 
         reasoning_steps = default(reasoning_steps, self.reasoning_steps)

HRM/hrm_with_act.py

@@ -89,6 +89,7 @@ def __init__(
     ):
         super().__init__()
         attn_layers_klass = Encoder if not causal else Decoder
+        self.causal = causal
 
         # input
 
@@ -169,12 +170,17 @@ def forward(
         detach_hiddens = True,
         one_step_grad = True,
         max_reasoning_steps = None,
-        adaptive_compute = None
+        adaptive_compute = None,
+        return_autoreg_loss = False
     ):
         batch, device = seq.shape[0], seq.device
 
         highest_hidden_index = self.num_networks - 1
 
+        if return_autoreg_loss:
+            assert self.causal
+            seq, labels = seq[:, :-1], seq[:, 1:]
+
         return_loss = exists(labels)
 
         # ACT related variables

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "HRM-pytorch"
-version = "0.1.4"
+version = "0.1.5"
 description = "The proposal from a Singaporean AGI company"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }