Fix weight params in conv1d and conv2d (#1245)

Author: Arjun31415

burn-core/src/nn/conv/conv1d.rs

@@ -50,8 +50,10 @@ pub struct Conv1dConfig {
 /// - bias:   Tensor of shape `[channels_out]`
 #[derive(Module, Debug)]
 pub struct Conv1d<B: Backend> {
-    weight: Param<Tensor<B, 3>>,
-    bias: Option<Param<Tensor<B, 1>>>,
+    /// Tensor of shape [channels_out, channels_in / groups, kernel_size]
+    pub weight: Param<Tensor<B, 3>>,
+    /// Tensor of shape `[channels_out]`
+    pub bias: Option<Param<Tensor<B, 1>>>,
     stride: usize,
     kernel_size: usize,
     dilation: usize,

burn-core/src/nn/conv/conv2d.rs

@@ -49,8 +49,10 @@ pub struct Conv2dConfig {
 /// - bias:   Tensor of shape `[channels_out]`
 #[derive(Module, Debug)]
 pub struct Conv2d<B: Backend> {
-    weight: Param<Tensor<B, 4>>,
-    bias: Option<Param<Tensor<B, 1>>>,
+    /// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]`
+    pub weight: Param<Tensor<B, 4>>,
+    /// Tensor of shape `[channels_out]`
+    pub bias: Option<Param<Tensor<B, 1>>>,
     stride: [usize; 2],
     kernel_size: [usize; 2],
     dilation: [usize; 2],

burn-core/src/nn/conv/conv_transpose1d.rs

@@ -51,8 +51,10 @@ pub struct ConvTranspose1dConfig {
 /// - bias:   Tensor of shape `[channels_out]`
 #[derive(Module, Debug)]
 pub struct ConvTranspose1d<B: Backend> {
-    weight: Param<Tensor<B, 3>>,
-    bias: Option<Param<Tensor<B, 1>>>,
+    /// Tensor of shape `[channels_in, channels_out / groups, kernel_size]`
+    pub weight: Param<Tensor<B, 3>>,
+    /// Tensor of shape `[channels_out]`
+    pub bias: Option<Param<Tensor<B, 1>>>,
     stride: usize,
     kernel_size: usize,
     dilation: usize,

burn-core/src/nn/conv/conv_transpose2d.rs

@@ -51,8 +51,10 @@ pub struct ConvTranspose2dConfig {
 /// - bias:   Tensor of shape `[channels_out]`
 #[derive(Module, Debug)]
 pub struct ConvTranspose2d<B: Backend> {
-    weight: Param<Tensor<B, 4>>,
-    bias: Option<Param<Tensor<B, 1>>>,
+    /// Tensor of shape `[channels_in, channels_out / groups, kernel_size_1, kernel_size_2]`
+    pub weight: Param<Tensor<B, 4>>,
+    /// Tensor of shape `[channels_out]`
+    pub bias: Option<Param<Tensor<B, 1>>>,
     stride: [usize; 2],
     kernel_size: [usize; 2],
     dilation: [usize; 2],