convert input to same type as weight for mixed precision training

When using CUDA AMP FMoE Linear throws type error as input is half() and weights are float()

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.float(), fwd_expert_count, self.weight, self.bias)
         return x
 
     def extra_repr(self) -> str: