Skip to content

Commit

Permalink
refactor: Unnest nn.Sequential for simpler handling
Browse files Browse the repository at this point in the history
The original hidden_layer in ResidualBlock consists of a nn.Linear and a nn.SiLU. Sperating them will affect nothing, but make the layer structure consistent with other child nodes output_layer and residual_layer
  • Loading branch information
TeddyHuang-00 committed Aug 29, 2024
1 parent 61fa1b2 commit 271aecf
Showing 1 changed file with 4 additions and 6 deletions.
10 changes: 4 additions & 6 deletions src/timesfm_torch/pytorch_patched_decoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,18 +260,16 @@ def __init__(
self.output_dims = output_dims

# Hidden Layer
self.hidden_layer = nn.Sequential(
nn.Linear(input_dims, hidden_dims),
nn.SiLU(),
)

self.hidden_layer = nn.Linear(input_dims, hidden_dims)
# Activation Function
self.act = nn.SiLU()
# Output Layer
self.output_layer = nn.Linear(hidden_dims, output_dims)
# Residual Layer
self.residual_layer = nn.Linear(input_dims, output_dims)

def forward(self, x):
hidden = self.hidden_layer(x)
hidden = self.act(self.hidden_layer(x))
output = self.output_layer(hidden)
residual = self.residual_layer(x)
return output + residual
Expand Down

0 comments on commit 271aecf

Please sign in to comment.