Skip to content

GitLab

Unverified Commit 3a0eb3e6 authored by Rick Ho's avatar Rick Ho Committed by GitHub
Browse files

Merge pull request #150 from santurini/patch-1 

convert input to same type as weight for mixed precision training
parents df715c9f 1f7c04c9
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
fmoe/linear.py
View file @ 3a0eb3e6
...@@ -63,14 +63,14 @@ class FMoELinear(nn.Module): ...@@ -63,14 +63,14 @@ class FMoELinear(nn.Module):
self.bias = nn.Parameter(torch.zeros(num_expert, out_feat)) self.bias = nn.Parameter(torch.zeros(num_expert, out_feat))
else: else:
self.register_parameter("bias", None) self.register_parameter("bias", None)
self.reset_parameters() self.reset_parameters()
def forward(self, inp, fwd_expert_count): def forward(self, inp, fwd_expert_count):
r""" r"""
Call MOE function 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 return x
def extra_repr(self) -> str: def extra_repr(self) -> str:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment