Cast input to weights type for AMP support

af8c4737 · Arturo Ghinassi · GitHub · a31fcc55 · af8c4737
Unverified Commit af8c4737 authored May 19, 2023 by Arturo Ghinassi 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: