mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-28 00:12:21 +08:00
Update moe.py
fixed not returning y in the forward method of MoEBlock
This commit is contained in:
parent
8ecad4cbb3
commit
b445dd3af3
@ -152,6 +152,8 @@ class MoEBlock(nn.Module):
|
|||||||
|
|
||||||
y = y + self.shared_experts(identity)
|
y = y + self.shared_experts(identity)
|
||||||
|
|
||||||
|
return y
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
||||||
|
|
||||||
@ -200,4 +202,4 @@ def test_moe():
|
|||||||
print(timing)
|
print(timing)
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
test_moe()
|
test_moe()
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user