oops

Author: lucidrains

HRM/hrm.py

@@ -171,7 +171,7 @@ def forward(
 
             is_last_step = index == (total_low_steps - 1)
 
-            context = torch.no_grad if one_step_grad and is_last_step else nullcontext
+            context = torch.no_grad if one_step_grad and not is_last_step else nullcontext
 
             with context():
                 # evaluate all networks depending on their period

pyproject.toml

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