diff --git a/src/neurons/leaky.rs b/src/neurons/leaky.rs
index 778738d..426a694 100644
--- a/src/neurons/leaky.rs
+++ b/src/neurons/leaky.rs
@@ -1,14 +1,24 @@
-use burn::{Tensor, module::Module, prelude::Backend};
+use burn::{Tensor, module::Module, prelude::Backend, tensor::backend::AutodiffBackend};
 
-#[derive(Debug, Module, Clone, Default)]
-pub struct LIF;
+#[derive(Debug, Module)]
+pub struct LIF<B: Backend> {
+    pub v: Tensor<B, 1>,
+}
 
-impl LIF {
-    pub fn new() -> Self {
-        Self
-    }
-
-    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
-        input
+impl<B: Backend> LIF<B> {
+    pub fn new(device: &B::Device) -> Self {
+        Self {
+            v: Tensor::zeros([1], device),
+        }
+    }
+}
+
+impl<B: AutodiffBackend> LIF<B> {
+    pub fn forward(&mut self, input: Tensor<B, 2>) -> Tensor<B, 2> {
+        input
+    }
+
+    pub fn backward(&mut self, x: Tensor<B, 2>) -> Tensor<B, 2> {
+        x
     }
 }