Merge pull request #68 from LukasHedegaard/develop

Fix state_index device after clean_state

Author: LukasHedegaard

CHANGELOG.md

@@ -8,6 +8,13 @@ From v1.0.0 and on, the project will adherence strictly to Semantic Versioning.
 
 ## Unpublished
 
+
+## [1.2.3] - 2023-06-16
+
+### Fixed
+- Ensure state_index remains on the same device after clean_state.
+
+
 ## [1.2.2] - 2023-05-24
 
 ### Fixed

continual/__about__.py

@@ -1,6 +1,6 @@
 import time
 
-__version__ = "1.2.2"
+__version__ = "1.2.3"
 __author__ = "Lukas Hedegaard"
 __author_email__ = "[email protected]"
 __license__ = "Apache-2.0"

continual/conv.py

@@ -155,7 +155,7 @@ def init_state(
         return (state_buffer, state_index, stride_index)
 
     def clean_state(self):
-        self.state_buffer = torch.tensor([])
+        self.state_buffer = torch.tensor([], device=self.state_buffer.device)
         self.state_index = torch.tensor(0)
         self.stride_index = torch.tensor(0)
 

continual/delay.py

@@ -76,7 +76,7 @@ def init_state(
         return state_buffer, state_index
 
     def clean_state(self):
-        self.state_buffer = torch.tensor([])
+        self.state_buffer = torch.tensor([], device=self.state_buffer.device)
         self.state_index = torch.tensor(0)
 
     def get_state(self):

continual/multihead_attention/retroactive_mha.py

@@ -276,11 +276,11 @@ def set_state(self, state: State):
         ) = state
 
     def clean_state(self):
-        self.d_mem = torch.tensor([])
-        self.AV_mem = torch.tensor([])
-        self.Q_mem = torch.tensor([])
-        self.K_T_mem = torch.tensor([])
-        self.V_mem = torch.tensor([])
+        self.d_mem = torch.tensor([], device=self.d_mem.device)
+        self.AV_mem = torch.tensor([], device=self.AV_mem.device)
+        self.Q_mem = torch.tensor([], device=self.Q_mem.device)
+        self.K_T_mem = torch.tensor([], device=self.K_T_mem.device)
+        self.V_mem = torch.tensor([], device=self.V_mem.device)
         self.stride_index = torch.tensor(0)
 
     def _forward_step(

continual/multihead_attention/single_output_mha.py

@@ -252,9 +252,9 @@ def set_state(self, state: State):
         ) = state
 
     def clean_state(self):
-        self.Q_mem = torch.tensor([])
-        self.K_T_mem = torch.tensor([])
-        self.V_mem = torch.tensor([])
+        self.Q_mem = torch.tensor([], device=self.Q_mem.device)
+        self.K_T_mem = torch.tensor([], device=self.K_T_mem.device)
+        self.V_mem = torch.tensor([], device=self.V_mem.device)
         self.stride_index = torch.tensor(0)
 
     @property

continual/pooling.py

@@ -166,7 +166,7 @@ def init_state(
         return state_buffer, state_index, stride_index
 
     def clean_state(self):
-        self.state_buffer = torch.tensor([])
+        self.state_buffer = torch.tensor([], device=self.state_buffer.device)
         self.state_index = torch.tensor(0)
         self.stride_index = torch.tensor(0)