Merge pull request #152 from santurini/patch-1

Cast input to weights type for AMP support

Merge pull request #152 from santurini/patch-1
Cast input to weights type for AMP support
50ac47f9 · Rick Ho · GitHub · a31fcc55 · af8c4737 · 50ac47f9
Unverified Commit 50ac47f9 authored May 19, 2023 by Rick Ho Committed by GitHub May 19, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

fmoe/linear.py fmoe/linear.py +1 -1

No files found.
--- a/fmoe/linear.py
+++ b/fmoe/linear.py
@@ -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: