don't save tensors, convert them to ints (#201)

jbondc · Oct 9, 2018 · d7d61f6 · d7d61f6
1 parent 39dd6b1
commit d7d61f6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/scripts/make_agent_demos.py b/scripts/make_agent_demos.py
@@ -15,6 +15,7 @@
 import time
 import numpy as np
 import blosc
+import torch
 
 import babyai.utils as utils
 
@@ -93,6 +94,8 @@ def generate_demos(n_episodes, valid, seed, shift=0):
         try:
             while not done:
                 action = agent.act(obs)['action']
+                if isinstance(action, torch.Tensor):
+                    action = action.item()
                 new_obs, reward, done, _ = env.step(action)
                 agent.analyze_feedback(reward, done)