use_gpu=False error fixed when conv is not used

carpedm20 · Jun 16, 2016 · 65ed340 · 65ed340
1 parent 95ed9d8
commit 65ed340
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/agents/experience.py b/agents/experience.py
@@ -49,7 +49,7 @@ def sample(self):
     rewards = self.rewards[indexes]
     terminals = self.terminals[indexes]
 
-    if self.data_format == 'NHWC':
+    if self.data_format == 'NHWC' and len(self.prestates.shape) == 4:
       return np.transpose(self.prestates, (0, 2, 3, 1)), actions, \
         rewards, np.transpose(self.poststates, (0, 2, 3, 1)), terminals
     else:

diff --git a/agents/history.py b/agents/history.py
@@ -13,7 +13,7 @@ def reset(self):
     self.history *= 0
 
   def get(self):
-    if self.data_format == 'NHWC':
+    if self.data_format == 'NHWC' and len(self.history.shape) == 3:
       return np.transpose(self.history, (1, 2, 0))
     else:
       return self.history