one more linear in mlp attention

lucidrains · lucidrains · commit 447768f13cba · 2022-11-26T16:53:19.000-08:00
diff --git a/equiformer_pytorch/equiformer_pytorch.py b/equiformer_pytorch/equiformer_pytorch.py
@@ -581,7 +581,8 @@ def __init__(
         # todo - needs a DTP here?
 
         self.to_values = nn.Sequential(
-            Gate(value_gate_fiber)
+            Gate(value_gate_fiber),
+            Linear(value_hidden_fiber, value_hidden_fiber)
         )
 
         # combining heads and projection out
diff --git a/equiformer_pytorch/version.py b/equiformer_pytorch/version.py
@@ -1,3 +1,3 @@
-__version__ = '0.0.16'
+__version__ = '0.0.17'
 
 __cuda_pkg_name__ = f'equiformer_pytorch_cuda_{__version__.replace(".", "_")}'

Original file line number	Diff line number	Diff line change
`@@ -581,7 +581,8 @@ def __init__(`
`581`	`581`	`# todo - needs a DTP here?`
`582`	`582`
`583`	`583`	`self.to_values = nn.Sequential(`
`584`		`- Gate(value_gate_fiber)`
	`584`	`+ Gate(value_gate_fiber),`
	`585`	`+ Linear(value_hidden_fiber, value_hidden_fiber)`
`585`	`586`	`)`
`586`	`587`
`587`	`588`	`# combining heads and projection out`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`		`-__version__ = '0.0.16'`
	`1`	`+__version__ = '0.0.17'`
`2`	`2`
`3`	`3`	`__cuda_pkg_name__ = f'equiformer_pytorch_cuda_{__version__.replace(".", "_")}'`