[Bugfix] Mark 'hidden_states' as mutable in moe_forward registration. (#20152)
Signed-off-by: Bill Nell <bnell@redhat.com>
This commit is contained in:
@ -1743,7 +1743,8 @@ def moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,
|
||||
direct_register_custom_op(
|
||||
op_name="moe_forward",
|
||||
op_func=moe_forward,
|
||||
mutates_args=[],
|
||||
mutates_args=["hidden_states"],
|
||||
fake_impl=moe_forward_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
tags=(torch.Tag.needs_fixed_stride_order, ),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user