[`NLLB-MoE`] Fix NLLB MoE 4bit inference (#27012)

fix NLLB MoE 4bit

[`NLLB-MoE`] Fix NLLB MoE 4bit inference (#27012)
fix NLLB MoE 4bit
244a53e0 · Younes Belkada · GitHub · cb45f71c · 244a53e0
Unverified Commit 244a53e0 authored Oct 23, 2023 by Younes Belkada Committed by GitHub Oct 23, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/models/nllb_moe/modeling_nllb_moe.py src/transformers/models/nllb_moe/modeling_nllb_moe.py +1 -1

No files found.
--- a/src/transformers/models/nllb_moe/modeling_nllb_moe.py
+++ b/src/transformers/models/nllb_moe/modeling_nllb_moe.py
@@ -417,7 +417,7 @@ class NllbMoeDenseActDense(nn.Module):
        if (
            isinstance(self.fc2.weight, torch.Tensor)
            and hidden_states.dtype != self.fc2.weight.dtype
-            and self.fc2.weight.dtype != torch.int8
+            and (self.fc2.weight.dtype != torch.int8 and self.fc2.weight.dtype != torch.uint8)
        ):
            hidden_states = hidden_states.to(self.fc2.weight.dtype)
        hidden_states = self.fc2(hidden_states)