def save_sample(step_dict, chunk_dir, chunk_item_idx):
    """
    Save a sample to the chunk directory.
    """
    # Save the json content
    time_stmp = time.time()
    while time.time() - time_stmp < 10.0:
        try:
            locks = []
            json_content = step_dict['json_content']
            file_path = os.path.join(chunk_dir, f"json_content_{chunk_item_idx}.json")
            lock = FileLock(file_path)
            locks.append(lock)
            lock.acquire_write_lock()
            with open(file_path, 'w') as file:
                json.dump(json_content, file, indent=4)
            lock.release_lock()
            # Save all other tensors in a npz
            file_path = os.path.join(chunk_dir, f"sample_{chunk_item_idx}.npz")
            lock = FileLock(file_path)
            locks.append(lock)
            lock.acquire_write_lock()
            with open(file_path, 'wb') as file:
                np.savez(
                    file,
                    step_id=step_dict['step_id'].numpy(),
                    state_chunk=step_dict['state_chunk'].numpy(),
                    state_chunk_time_mask=step_dict['state_chunk_time_mask'].numpy(),
                    action_chunk=step_dict['action_chunk'].numpy(),
                    action_chunk_time_mask=step_dict['action_chunk_time_mask'].numpy(),
                    state_vec_mask=step_dict['state_vec_mask'].numpy(),
                    past_frames_0=step_dict['past_frames_0'].numpy(),
                    past_frames_0_time_mask=step_dict['past_frames_0_time_mask'].numpy(),
                    past_frames_1=step_dict['past_frames_1'].numpy(),
                    past_frames_1_time_mask=step_dict['past_frames_1_time_mask'].numpy(),
                    past_frames_2=step_dict['past_frames_2'].numpy(),
                    past_frames_2_time_mask=step_dict['past_frames_2_time_mask'].numpy(),
                    past_frames_3=step_dict['past_frames_3'].numpy(),
                    past_frames_3_time_mask=step_dict['past_frames_3_time_mask'].numpy(),
                    state_std=step_dict['state_std'].numpy(),
                    state_mean=step_dict['state_mean'].numpy(),
                    state_norm=step_dict['state_norm'].numpy(),            
                )
            lock.release_lock()
            return
        except KeyboardInterrupt:
            for lock in locks:
                lock.release_lock()
            raise KeyboardInterrupt
        except BaseException:
            for lock in locks:
                lock.release_lock()
            continue
    # raise RuntimeError("Failed to save sample.")
    print("Failed to save sample.")

    Save a sample to the chunk directory.
    
policy/RDT/data/producer.py