update

algorithms/worldmem/df_video.py

@@ -804,9 +804,12 @@ class WorldMemMinecraft(DiffusionForcingBase):
             new_c2w_mat = euler_to_camera_to_world_matrix(first_pose)
             self_memory_c2w = new_c2w_mat[None, None].to(device)
             self_frame_idx = torch.tensor([[context_frames_idx]]).to(device)
-            return first_frame, self_frames, self_poses, self_memory_c2w, self_frame_idx
+            return first_frame.cpu(), self_frames.cpu(), self_poses.cpu(), self_memory_c2w.cpu(), self_frame_idx.cpu()
         else:
             last_frame = self_frames[-1].clone()
+            self_poses = self_poses.to(device)
+            self_memory_c2w = self_memory_c2w.to(device)
+            self_frame_idx = self_frame_idx.to(device)
             last_pose_condition = self_poses[-1].clone()
             last_pose_condition[:,3:] = last_pose_condition[:,3:] // 15
             new_pose_condition_offset = self.pose_prediction_model(last_frame.to(device), curr_actions[None].to(device), last_pose_condition)
@@ -900,7 +903,7 @@ class WorldMemMinecraft(DiffusionForcingBase):
 
         xs_pred = self.decode(xs_pred[n_context_frames:].to(device)).cpu()
 
-        return xs_pred[-1,0], self_frames, self_poses.cpu(), self_memory_c2w.cpu(), self_frame_idx.cpu()
+        return xs_pred[-1,0].cpu(), self_frames.cpu(), self_poses.cpu(), self_memory_c2w.cpu(), self_frame_idx.cpu()
 
 
     def reset(self):