From f5226b5152e8b61b2e342087512d4a01e0f11872 Mon Sep 17 00:00:00 2001
From: Wenwen Qu <vinny_qu@163.com>
Date: Mon, 8 Jan 2024 16:23:53 +0800
Subject: [PATCH] refactor code

---
 internlm/moe/base_moe.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internlm/moe/base_moe.py b/internlm/moe/base_moe.py
index 5f19e07..c97ec83 100644
--- a/internlm/moe/base_moe.py
+++ b/internlm/moe/base_moe.py
@@ -31,5 +31,5 @@ class BaseMoELayer(Base):
         self.ep_group = ep_group
         self.ep_size = ep_size
         self.num_local_experts = num_local_experts
-        self.l_aux = torch.zeros(1, device=torch.cuda.current_device())
+        self.l_aux = torch.tensor(0.0, device=torch.cuda.current_device(), dtype=gpc.config.model.get("dtype"))
         self.exp_counts = None