Skip to content

GitLab

Unverified Commit af8c4737 authored by Arturo Ghinassi's avatar Arturo Ghinassi Committed by GitHub
Browse files

Cast input to weights type for AMP support

parent a31fcc55
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
fmoe/linear.py
View file @ af8c4737
...@@ -70,7 +70,7 @@ class FMoELinear(nn.Module): ...@@ -70,7 +70,7 @@ class FMoELinear(nn.Module):
r""" r"""
Call MOE function 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 return x
def extra_repr(self) -> str: 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