Merge pull request 'gates where not saved as well' (#25) from bug_fixes into main
Reviewed-on: #25
This commit is contained in:
commit
2eee873247
|
@ -193,8 +193,8 @@ class AIIAmoe(AIIA):
|
||||||
|
|
||||||
# To generate gating weights, we first need to determine the feature dimension.
|
# 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,
|
# 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).
|
# we obtain a tensor of shape (B, C) where C is the number of channels (here assumed to be 224).
|
||||||
gate_in_features = 410 # Adjust this if your expert output changes.
|
gate_in_features = 512 # Adjust this if your expert output changes.
|
||||||
|
|
||||||
# Create a gating network that maps the aggregated features to num_experts weights.
|
# Create a gating network that maps the aggregated features to num_experts weights.
|
||||||
self.gate = nn.Sequential(
|
self.gate = nn.Sequential(
|
||||||
|
|
Loading…
Reference in New Issue