chore: simplify spike handling in SNNTorch models and reset hidden states in MNIST model

Author: Khushiyant

benchmarks/bechmark_sota.py

@@ -222,16 +222,16 @@ def forward(self, x):
             x_t = x[:, t] if x.shape[1] == self.n_steps else x[t]
 
             cur = self.conv1(x_t)
-            spk, _ = self.lif1(cur)
+            spk = self.lif1(cur)
             cur = self.pool1(spk)
 
             cur = self.conv2(cur)
-            spk, _ = self.lif2(cur)
+            spk = self.lif2(cur)
             cur = self.pool2(spk)
 
             cur = self.flatten(cur)
             cur = self.fc1(cur)
-            spk, _ = self.lif3(cur)
+            spk = self.lif3(cur)
             cur = self.fc2(spk)
 
             outputs.append(cur)
@@ -261,6 +261,7 @@ def __init__(self, n_steps=10, beta=0.9):
         self.fc2 = nn.Linear(128, 10)
 
     def forward(self, x):
+        # Reset hidden states
         self.lif1.init_leaky()
         self.lif2.init_leaky()
         self.lif3.init_leaky()
@@ -270,16 +271,16 @@ def forward(self, x):
             x_t = x[:, t] if x.shape[1] == self.n_steps else x[t]
 
             cur = self.conv1(x_t)
-            spk, _ = self.lif1(cur)
+            spk = self.lif1(cur)
             cur = self.pool1(spk)
 
             cur = self.conv2(cur)
-            spk, _ = self.lif2(cur)
+            spk = self.lif2(cur)
             cur = self.pool2(spk)
 
             cur = self.flatten(cur)
             cur = self.fc1(cur)
-            spk, _ = self.lif3(cur)
+            spk = self.lif3(cur)
             cur = self.fc2(spk)
 
             outputs.append(cur)

pyproject.toml

@@ -17,14 +17,21 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [project.optional-dependencies]
-dev = ["pytest>=9.0.2", "pytest-cov>=7.0.0"]
+dev = [
+    "pytest>=9.0.2",
+    "pytest-cov>=7.0.0",
+    "psutil>=7.2.1",
+    "snntorch>=0.9.4",
+    "gputil>=1.4.0",
+    "spikingjelly>=0.0.0.0.14",
+]
 docs = [
     "sphinx>=8.1.3",
     "sphinx-autodoc-typehints>=3.0.1",
     "sphinx-rtd-theme>=3.0.0",
     "myst-parser>=4.0.1",
     "furo>=2025.9.25",
-    "ruff>=0.14.10"
+    "ruff>=0.14.10",
 ]
 
 [tool.hatch.build.targets.wheel]
@@ -37,3 +44,6 @@ branch = "main"
 upload_to_pypi = false
 upload_to_release = true
 build_command = "pip install build && python -m build"
+
+[dependency-groups]
+dev = []