Skip to content

GitLab

Commit 3d8e8a43 authored by Jiezhong Qiu's avatar Jiezhong Qiu
Browse files

follow pylint to remove lambda

parent e36c3c4c
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 2 deletions
fmoe/layers.py
View file @ 3d8e8a43
...@@ -191,8 +191,7 @@ class FMoE(nn.Module): ...@@ -191,8 +191,7 @@ class FMoE(nn.Module):
gate_top_k_idx, gate_score = self.gate(inp) gate_top_k_idx, gate_score = self.gate(inp)
# to: (BxLxtop_k) x d_model # to: (BxLxtop_k) x d_model
inp = inp.repeat_interleave(repeats=self.top_k, dim=0) inp = inp.repeat_interleave(repeats=self.top_k, dim=0)
expert_fn = lambda inp, fec: self.expert_fn(inp, fec) x = _fmoe_general_global_forward(inp, gate_top_k_idx, self.expert_fn,
x = _fmoe_general_global_forward(inp, gate_top_k_idx, expert_fn,
self.num_expert, self.world_size) self.num_expert, self.world_size)
# to: (BxL) x top_k x d_model # to: (BxL) x top_k x d_model
x = x.view(-1, self.top_k, self.d_model) x = x.view(-1, self.top_k, self.d_model)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment