<jupyter_start><jupyter_text>Now for the cpu:<jupyter_code># Running parallel sklearn
# cpu version doesn't work with parallel laoding for some reason
start = time()
dataset = VideoDataset(movie_path, parallel=False)
loader = DataLoader(dataset, batch_size=None, pin_memory=True)

preprocessor_ini = preprocessing_blob(mask)
localize_ini = localize_blob(default_blob_detector_params())
initial_position, initial_frame = _initialize(
    loader, preprocessor_ini, localize_ini, 100
)

preprocessor_main = preprocessing_kmeans(mask, device="cpu")
localize_main = localize_kmeans_sklearn(120, 1e-4)
locs_cpu = _localize(
    loader, preprocessor_main, localize_main, initial_position, n_frames, "cpu"
)
stop = time()
print(f"Time for CPU run: {stop - start}s")

np.max(np.stack(locs_cpu) - torch.stack(locs_gpu).cpu().numpy())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10902696365767878, 'fraction_numerical': 0.009378663540445486, 'mean_word_length': 7.989473684210527, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '360442', 'license': 'MIT', 'n_tokens_mistral': 302, 'n_tokens_neox': 292, 'n_words': 61, 'path': 'tests/compare_gpu_cpu.ipynb', 'repo_name': 'nicolaseberle/flyTracker', 'size': '853'}