update doc

hpcaitech · Aug 14, 2023 · 00cb06e · 00cb06e
1 parent c84f37b
commit 00cb06e
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/colossalai/nn/layer/moe/moe_param.py b/colossalai/nn/layer/moe/moe_param.py
@@ -14,5 +14,13 @@ def is_moe_param(tensor: torch.Tensor) -> bool:
     return hasattr(tensor, "moe_info")
 
 
-def set_moe_param_info(tensor: torch.Tensor, moe_info: dict):
+def set_moe_param_info(tensor: torch.Tensor, moe_info: dict) -> None:
+    """
+    Set moe info for the given tensor.
+
+    Args:
+        tensor (torch.Tensor): The tensor to be set.
+        moe_info (dict): The moe info to be set.
+
+    """
     tensor.__setattr__('moe_info', moe_info)