espnet2.sds.end_to_end.mini_omni.litgpt.model.LLaMAMoE

source

class espnet2.sds.end_to_end.mini_omni.litgpt.model.LLaMAMoE(config: Config)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Forward method derived from:

https://github.com/mistralai/mistral-src/blob/b46d6/ moe_one_file_ref.py#L203-L219 See also figure 1 in https://arxiv.org/abs/2211.15841