Skip to content

GitLab

Unverified Commit 50ac47f9 authored by Rick Ho's avatar Rick Ho Committed by GitHub
Browse files

Merge pull request #152 from santurini/patch-1 

Cast input to weights type for AMP support
parents a31fcc55 af8c4737
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
fmoe/linear.py
View file @ 50ac47f9
......@@ -70,7 +70,7 @@ class FMoELinear(nn.Module):
r"""
Call MOE function
"""
x = MOELinear.apply(inp, fwd_expert_count, self.weight, self.bias)
x = MOELinear.apply(inp.type_as(self.weight), fwd_expert_count, self.weight, self.bias)
return x
def extra_repr(self) -> str:
......
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