address #30

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.6.4"
+version = "0.6.5"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

transfusion_pytorch/transfusion.py

@@ -118,6 +118,11 @@ def default(v, d):
 def identity(t):
     return t
 
+def always(val):
+    def inner(*args, **kwargs):
+        return val
+    return inner
+
 def first(it):
     return it[0]
 
@@ -776,7 +781,7 @@ def __init__(
             nn.Linear(dim, heads),
             nn.Sigmoid(),
             Rearrange('b n h -> b h n 1') # add head dimension
-        ) if learned_value_residual_mix else (lambda _: 0.5)
+        ) if learned_value_residual_mix else always(0.5)
 
         self.to_gates = nn.Sequential(
             nn.Linear(dim, heads, bias = False),