add dropout to input

src/main.rs

@@ -1,6 +1,6 @@
 use burn::backend::Autodiff;
 use burn::nn::loss::{MseLoss, Reduction};
-use burn::nn::{Linear, LinearConfig, Relu};
+use burn::nn::{Dropout, DropoutConfig, Linear, LinearConfig, Relu};
 use burn::optim::{AdamConfig, GradientsParams, Optimizer};
 use burn::prelude::*;
 
@@ -30,6 +30,7 @@ const STRIDE_SIZE: u32 = (SIZE.x * size_of::<f32>() as u32).next_multiple_of(256
 #[derive(Module, Debug)]
 struct Net<B: Backend> {
 	input: Linear<B>,
+	dropout: Dropout,
 	hidden: [Linear<B>; HIDDEN.len() - 1],
 	output: Linear<B>,
 	activation: Relu,
@@ -46,8 +47,10 @@ impl<B: Backend> Net<B> {
 			layer
 		});
 		let output = LinearConfig::new(last_size, OUTPUT).init(device);
+		let dropout = DropoutConfig::new(0.1).init();
 		Self {
 			input,
+			dropout,
 			hidden,
 			output,
 			activation: Relu::new(),
@@ -55,6 +58,7 @@ impl<B: Backend> Net<B> {
 	}
 	fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
 		let x = self.input.forward(input);
+		let x = self.dropout.forward(x);
 		let mut x = self.activation.forward(x);
 		for layer in &self.hidden {
 			x = layer.forward(x);