From b01e3907a16efcd4972a1bfa1ff0691363a2433b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 10 Dec 2023 17:21:49 -0800 Subject: [PATCH] mixtral touch up: two lines --- examples/mixtral.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/mixtral.py b/examples/mixtral.py index f5a3f35ec4..8593cf4a4e 100644 --- a/examples/mixtral.py +++ b/examples/mixtral.py @@ -16,7 +16,8 @@ class MixtureFeedForward: top = sorted(enumerate(choice), key=lambda x: -x[1]) norm = top[0][1] + top[1][1] e1, e2 = self.experts[top[0][0]], self.experts[top[1][0]] - ret = e1(x.to(e1.w1.weight.device)).to(x.device) * (top[0][1]/norm) + e2(x.to(e2.w1.weight.device)).to(x.device) * (top[1][1]/norm) + ret = e1(x.to(e1.w1.weight.device)).to(x.device) * (top[0][1]/norm) + \ + e2(x.to(e2.w1.weight.device)).to(x.device) * (top[1][1]/norm) return ret if __name__ == "__main__":