bug fix

tracel-ai · May 3, 2024 · e63fea4 · e63fea4
1 parent 8d226c7
commit e63fea4
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 19 deletions.
diff --git a/crates/burn-import/src/burn/node/prelu.rs b/crates/burn-import/src/burn/node/prelu.rs
@@ -89,7 +89,6 @@ impl<PS: PrecisionSettings> NodeCodegen<PS> for PReluNode<PS> {
         }
     }
     fn register_imports(&self, imports: &mut BurnImports) {
-        imports.register("burn::nn::PRelu");
         imports.register("burn::nn::prelu::PRelu");
         imports.register("burn::nn::prelu::PReluConfig");
     }

diff --git a/crates/burn-import/src/onnx/op_configuration.rs b/crates/burn-import/src/onnx/op_configuration.rs
@@ -1,7 +1,7 @@
 use burn::nn::{
     conv::{Conv1dConfig, Conv2dConfig, ConvTranspose2dConfig},
     pool::{AvgPool2dConfig, MaxPool2dConfig},
-    BatchNormConfig, DropoutConfig, LayerNormConfig, LinearConfig, PReluConfig, PaddingConfig1d,
+    BatchNormConfig, DropoutConfig, LayerNormConfig, LinearConfig, PaddingConfig1d,
     PaddingConfig2d,
 };
 
@@ -120,22 +120,6 @@ pub fn max_pool2d_config(curr: &Node) -> MaxPool2dConfig {
         .with_padding(padding)
         .with_dilation([dilations[0] as usize, dilations[1] as usize])
 }
-pub fn prelu_config(curr: &Node) -> PReluConfig {
-    let mut alpha = 0.01;
-    let mut num_parameters = 0;
-    for (key, value) in curr.attrs.iter() {
-        match key.as_str() {
-            "alpha" => alpha = value.clone().into_f32(),
-            "num_parameters" => num_parameters = value.clone().into_i32(),
-            _ => {}
-        }
-    }
-
-    PReluConfig::new()
-        .with_num_parameters(num_parameters as usize)
-        .with_alpha(alpha as f64)
-}
-
 pub fn conv_transpose2d_config(curr: &Node) -> ConvTranspose2dConfig {
     let mut attrs = curr.attrs.clone();
     let kernel_shape = attrs

diff --git a/crates/burn-import/src/onnx/to_burn.rs b/crates/burn-import/src/onnx/to_burn.rs
@@ -5,6 +5,7 @@ use std::{
 };
 
 use burn::{
+    nn::PReluConfig,
     record::{FullPrecisionSettings, HalfPrecisionSettings, PrecisionSettings},
     tensor::{DataSerialize, Element},
 };
@@ -701,7 +702,7 @@ impl OnnxGraph {
         let input = node.inputs.first().unwrap().to_tensor_type();
         let output = node.outputs.first().unwrap().to_tensor_type();
         let weight = extract_data_serialize::<PS::FloatElem>(1, &node).unwrap();
-        let config = prelu_config(&node);
+        let config = PReluConfig::new();
         let name = &node.name;
         PReluNode::<PS>::new(name, input, output, weight, config)
     }