fixed broken animation

mfinzi · Apr 15, 2021 · dfa9c03 · dfa9c03
1 parent 8bf0ae5
commit dfa9c03
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/experiments/trainer/hamiltonian_dynamics.py b/experiments/trainer/hamiltonian_dynamics.py
@@ -111,7 +111,7 @@ def sample_initial_conditions(self,bs):
         raise NotImplementedError
     def animate(self, zt=None):
         if zt is None:
-            zt = np.asarray(self.integrate(self.sample_initial_conditions(10),self.T_long))
+            zt = np.asarray(self.integrate(self.sample_initial_conditions(10)[0],self.T_long))
         # bs, T, 2nd
         if len(zt.shape) == 3:
             j = np.random.randint(zt.shape[0])
@@ -453,4 +453,4 @@ def __call__(self,cfg,i=None):
             if self.strict: raise
             outcome = e
         del trainer
-        return cfg, outcome
+        return cfg, outcome