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

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

fmoe/linear.py

@@ -63,14 +63,14 @@ class FMoELinear(nn.Module):
             self.bias = nn.Parameter(torch.zeros(num_expert, out_feat))
         else:
             self.register_parameter("bias", None)
-          
+
         self.reset_parameters()
 
     def forward(self, inp, fwd_expert_count):
         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: