fixed store_device issue (#249)

Co-authored-by: Chris Nota <[email protected]>

Author: benblack769

all/bodies/vision.py

@@ -69,10 +69,9 @@ def update(self, key, value):
         x = {}
         for k in self.keys():
             if not k == key:
-                x[k] = super().__getitem__(k)
+                x[k] = dict.__getitem__(self, k)
         x[key] = value
-        state = LazyState(x, device=self.device)
-        state.to_cache = self.to_cache
+        state = LazyState.from_state(x, x['observation'], self.to_cache)
         return state
 
     def to(self, device):

all/memory/replay_buffer.py

@@ -99,7 +99,7 @@ def store(self, state, action, next_state):
         if state is None or state.done:
             return
         idx = self.pos
-        super()._add((state, action, next_state))
+        super().store(state, action, next_state)
         self._it_sum[idx] = self._max_priority ** self._alpha
         self._it_min[idx] = self._max_priority ** self._alpha