rxnn 0.2.4__tar.gz → 0.2.5__tar.gz

{rxnn-0.2.4 → rxnn-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.4
+Version: 0.2.5
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.4 → rxnn-0.2.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.4"
+version = "0.2.5"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.4 → rxnn-0.2.5}/src/rxnn/memory/stm.py

@@ -48,22 +48,23 @@ class ShortTermMemory(nn.Module):
         if not self.is_trainable:
             self.is_trainable = True
             initial_stm = self.memory.clone()
-            del self.memory
+            delattr(self, 'memory')
             self.memory = nn.Parameter(initial_stm)
 
     def freeze(self):
         if self.is_trainable:
             self.requires_grad_(False)
             trained_stm = self.memory.clone()
-            del self.memory
+            delattr(self, 'memory')
             self.register_buffer('memory', trained_stm)
 
     def reset(self, init_type: str = None):
-        self.memory = self._init_tensor(init_type)
+        self.memory.copy_(self._init_tensor(init_type))
 
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
-        self.memory = self._init_tensor(init_type)
+        delattr(self, 'memory')
+        self.register_buffer('memory', self._init_tensor(init_type))
 
     def batched_memory(self, batch_size: int, init_type: str = None):
         if init_type is not None:
@@ -71,7 +72,8 @@ class ShortTermMemory(nn.Module):
                 'STM init type must be one of "normal", "standard", "uniform", "ones", "zeros"'
             self.init_type = init_type
         self.batch_size = batch_size
-        self.memory = self._init_tensor()
+        delattr(self, 'memory')
+        self.register_buffer('memory', self._init_tensor())
 
     def single_memory(self, init_type: str = None, use_mean_from_batch: bool = False):
         if init_type is not None:
@@ -81,7 +83,9 @@ class ShortTermMemory(nn.Module):
         self.batch_size = 1
         if use_mean_from_batch:
             batch_mean = self.memory.mean(dim=(1, 2, 3), keepdim=True)
-            self.memory = self._init_tensor()
+            delattr(self, 'memory')
+            self.register_buffer('memory', self._init_tensor())
             self.memory.copy_(batch_mean)
         else:
-            self.memory = self._init_tensor()
+            delattr(self, 'memory')
+            self.register_buffer('memory', self._init_tensor())