Merge pull request #151 from laekov/revert-150-patch-1

Revert "convert input to same type as weight for mixed precision training"

Merge pull request #151 from laekov/revert-150-patch-1
Revert "convert input to same type as weight for mixed precision training"
a31fcc55 · Rick Ho · GitHub · 3a0eb3e6 · f1a6e6fd · a31fcc55
Unverified Commit a31fcc55 authored May 19, 2023 by Rick Ho Committed by GitHub May 19, 2023
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

fmoe/linear.py fmoe/linear.py +2 -2

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.type_as(self.weight), fwd_expert_count, self.weight, self.bias)
+        x = MOELinear.apply(inp, fwd_expert_count, self.weight, self.bias)
        return x

    def extra_repr(self) -> str: