Unverified Commit c6c98305 authored by bnellnm's avatar bnellnm Committed by GitHub
Browse files

[Bugfix] Mark 'hidden_states' as mutable in moe_forward registration. (#20152)


Signed-off-by: default avatarBill Nell <bnell@redhat.com>
parent aafabaa0
......@@ -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, ),
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment