refactoring

NeuroBench · Jul 24, 2024 · 61eee43 · 61eee43
1 parent 0d2f25b
commit 61eee43
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
@@ -25,9 +25,7 @@
 
 
 class TestStaticMetrics(unittest.TestCase):
-
     def setUp(self):
-
         self.dummy_net = nn.Module()
 
         self.net = models.net
@@ -353,13 +351,16 @@ def test_synaptic_ops(self):
 
         syn = synaptic_operations()
         syn_ops = syn(model, out, inp)
-        
+
         self.assertEqual(syn_ops["Effective_MACs"], 615)
         self.assertEqual(syn_ops["Effective_ACs"], 0)
 
         # test RNN network
         batch_size = 2
-        inp = [torch.ones(batch_size, 25), torch.ones(batch_size, 5)]  # input, (hidden, cell)
+        inp = [
+            torch.ones(batch_size, 25),
+            torch.ones(batch_size, 5),
+        ]  # input, (hidden, cell)
         inp[0][0, 0] = 4  # avoid getting classified as snn
         model = TorchModel(self.net_RNN)
 
@@ -375,7 +376,10 @@ def test_synaptic_ops(self):
 
         # test GRU network
         batch_size = 2
-        inp = [torch.ones(batch_size, 25), torch.ones(batch_size, 5)]  # input, (hidden, cell)
+        inp = [
+            torch.ones(batch_size, 25),
+            torch.ones(batch_size, 5),
+        ]  # input, (hidden, cell)
         inp[0][0, 0] = 4  # avoid getting classified as snn
         model = TorchModel(self.net_GRU)
 
@@ -403,6 +407,7 @@ def test_membrane_potential_updates(self):
 
         self.assertEqual(tot_mem_updates, 50)
 
+
 # TODO: refactor this metric if needed
 # def test_neuron_update_metric():
 #     net_relu_0 = nn.Sequential(