selfattention block: Remove the fc linear layer if it is not used

Signed-off-by: John Zielke <[email protected]>
Project-MONAI · Feb 4, 2025 · 547ac94 · 547ac94
1 parent 8dcb9dc
commit 547ac94
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/monai/networks/blocks/selfattention.py b/monai/networks/blocks/selfattention.py
@@ -106,7 +106,11 @@ def __init__(
 
         self.num_heads = num_heads
         self.hidden_input_size = hidden_input_size if hidden_input_size else hidden_size
-        self.out_proj = nn.Linear(self.inner_dim, self.hidden_input_size)
+        self.out_proj: Union[nn.Linear, nn.Identity]
+        if include_fc:
+            self.out_proj = nn.Linear(self.inner_dim, self.hidden_input_size)
+        else:
+            self.out_proj = nn.Identity()
 
         self.qkv: Union[nn.Linear, nn.Identity]
         self.to_q: Union[nn.Linear, nn.Identity]