address #10

q_transformer/agent.py

@@ -206,7 +206,6 @@ def __init__(
         if condition_on_text:
             text_embeds_path = mem_path / TEXT_EMBEDS_FILENAME
             text_embed_shape = environment.text_embed_shape
-
             self.text_embed_shape = text_embed_shape
             self.text_embeds = open_memmap(str(text_embeds_path), dtype = 'float32', mode = 'w+', shape = (*prec_shape, *text_embed_shape))
 
@@ -269,11 +268,21 @@ def forward(self):
                 curr_state = next_state
 
             if self.condition_on_text:
-                del self.text_embeds
+                self.text_embeds.flush()
+
+            self.states.flush()
+            self.actions.flush()
+            self.rewards.flush()
+            self.dones.flush()
+
+        # close memmap
+
+        if self.condition_on_text:
+            del self.text_embeds
 
-            del self.states
-            del self.actions
-            del self.rewards
-            del self.dones
+        del self.states
+        del self.actions
+        del self.rewards
+        del self.dones
 
         print(f'completed, memories stored to {self.memories_dataset_folder.resolve()}')

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'q-transformer',
   packages = find_packages(exclude=[]),
-  version = '0.1.12',
+  version = '0.1.14',
   license='MIT',
   description = 'Q-Transformer',
   author = 'Phil Wang',