diff --git a/rust/tvm/src/ir/relay/attrs/nn.rs b/rust/tvm/src/ir/relay/attrs/nn.rs index f743534e5f61..cb96f0fbf588 100644 --- a/rust/tvm/src/ir/relay/attrs/nn.rs +++ b/rust/tvm/src/ir/relay/attrs/nn.rs @@ -94,3 +94,15 @@ pub struct SoftmaxAttrsNode { pub base: BaseAttrsNode, pub axis: i32, } + +#[repr(C)] +#[derive(Object)] +#[ref_name = "BatchNormAttrs"] +#[type_key = "relay.attrs.BatchNormAttrs"] +pub struct BatchNormAttrsNode { + pub base: BaseAttrsNode, + pub axis: i32, + pub epsilon: f64, + pub center: bool, + pub scale: bool, +}