gates where not saved as well

This commit is contained in:
Falko Victor Habel 2025-03-03 19:01:29 +01:00
parent 91728a0469
commit 35edadf727
1 changed files with 2 additions and 2 deletions

View File

@ -193,8 +193,8 @@ class AIIAmoe(AIIA):
# To generate gating weights, we first need to determine the feature dimension.
# Each expert is assumed to return an output of shape (B, C, H, W); after averaging over H and W,
# we obtain a tensor of shape (B, C) where C is the number of channels (here assumed to be 410).
gate_in_features = 410 # Adjust this if your expert output changes.
# we obtain a tensor of shape (B, C) where C is the number of channels (here assumed to be 224).
gate_in_features = 512 # Adjust this if your expert output changes.
# Create a gating network that maps the aggregated features to num_experts weights.
self.gate = nn.Sequential(