Don't register mem as buffer

jeshraghian · Apr 17, 2024 · 67973ed · 67973ed
1 parent bdc1b49
commit 67973ed
Show file tree

Hide file tree

Showing 8 changed files with 16 additions and 34 deletions.
diff --git a/snntorch/_neurons/alpha.py b/snntorch/_neurons/alpha.py
@@ -129,13 +129,9 @@ def __init__(
             self.state_function = self._base_int
 
     def _init_mem(self):
-        syn_exc = torch.zeros(1)
-        syn_inh = torch.zeros(1)
-        mem = torch.zeros(1)
-
-        self.register_buffer("syn_exc", syn_exc)
-        self.register_buffer("syn_inh", syn_inh)
-        self.register_buffer("mem", mem)
+        self.syn_exc = torch.zeros(1)
+        self.syn_inh = torch.zeros(1)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.syn_exc = torch.zeros_like(

diff --git a/snntorch/_neurons/lapicque.py b/snntorch/_neurons/lapicque.py
@@ -226,8 +226,7 @@ def __init__(
             self.state_function = self._base_int
 
     def _init_mem(self):
-        mem = torch.zeros(1)
-        self.register_buffer("mem", mem)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.mem = torch.zeros_like(self.mem, device=self.mem.device)

diff --git a/snntorch/_neurons/leaky.py b/snntorch/_neurons/leaky.py
@@ -182,8 +182,7 @@ def __init__(
         self.reset_delay = reset_delay
 
     def _init_mem(self):
-        mem = torch.zeros(1)
-        self.register_buffer("mem", mem)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.mem = torch.zeros_like(self.mem, device=self.mem.device)

diff --git a/snntorch/_neurons/rleaky.py b/snntorch/_neurons/rleaky.py
@@ -292,10 +292,8 @@ def __init__(
         self.reset_delay = reset_delay
 
     def _init_mem(self):
-        spk = torch.zeros(1)
-        mem = torch.zeros(1)
-        self.register_buffer("spk", spk)
-        self.register_buffer("mem", mem)
+        self.spk = torch.zeros(1)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.spk = torch.zeros_like(self.spk, device=self.spk.device)

diff --git a/snntorch/_neurons/rsynaptic.py b/snntorch/_neurons/rsynaptic.py
@@ -307,13 +307,9 @@ def __init__(
         self.reset_delay = reset_delay
 
     def _init_mem(self):
-        spk = torch.zeros(1)
-        syn = torch.zeros(1)
-        mem = torch.zeros(1)
-
-        self.register_buffer("spk", spk)
-        self.register_buffer("syn", syn)
-        self.register_buffer("mem", mem)
+        self.spk = torch.zeros(1)
+        self.syn = torch.zeros(1)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.spk = torch.zeros_like(self.spk, device=self.spk.device)

diff --git a/snntorch/_neurons/sconv2dlstm.py b/snntorch/_neurons/sconv2dlstm.py
@@ -273,10 +273,8 @@ def __init__(
         )
 
     def _init_mem(self):
-        syn = torch.zeros(1)
-        mem = torch.zeros(1)
-        self.register_buffer("syn", syn)
-        self.register_buffer("mem", mem)
+        self.syn = torch.zeros(1)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.syn = torch.zeros_like(self.syn, device=self.syn.device)

diff --git a/snntorch/_neurons/slstm.py b/snntorch/_neurons/slstm.py
@@ -203,10 +203,8 @@ def __init__(
         )
 
     def _init_mem(self):
-        syn = torch.zeros(1)
-        mem = torch.zeros(1)
-        self.register_buffer("syn", syn)
-        self.register_buffer("mem", mem)
+        self.syn = torch.zeros(1)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.syn = torch.zeros_like(self.syn, device=self.syn.device)

diff --git a/snntorch/_neurons/synaptic.py b/snntorch/_neurons/synaptic.py
@@ -198,10 +198,8 @@ def __init__(
         self.reset_delay = reset_delay
 
     def _init_mem(self):
-        syn = torch.zeros(1)
-        mem = torch.zeros(1)
-        self.register_buffer("syn", syn)
-        self.register_buffer("mem", mem)
+        self.syn = torch.zeros(1)
+        self.mem = torch.zeros(1)
 
     def reset_mem(self):
         self.syn = torch.zeros_like(self.syn, device=self.syn.device)