Segmentation Fault in flow_predictor_raft.py script

[hadillivvy]

Hello! I've been facing a segmentation fault error right after running the flowmap.overfit script and it seems to be from this block of code in flow_predictor_raft.py (line 41):

flow = [
            self.raft(
                source_chunk * 2 - 1,
                target_chunk * 2 - 1,
                num_flow_updates=self.cfg.num_flow_updates,
            )[-1]
            for source_chunk, target_chunk in zip(
                bar(source.split(self.cfg.max_batch_size)),
                target.split(self.cfg.max_batch_size),
            )
        ]

Here is the stack trace:
python -m flowmap.overfit dataset=images dataset.images.root=/home/hadil/Work/flowmap/data
rm: cannot remove 'outputs/local': No such file or directory
Precomputing optical flow.
Computing RAFT flow: 0%| | 0/1 [00:00<?, ?it/s]Fatal Python error: Segmentation fault

Thread 0x00007e39449e3640 (most recent call first):
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/threading.py", line 331 in wait
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/threading.py", line 629 in wait
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/tqdm/_monitor.py", line 60 in run
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/threading.py", line 1045 in _bootstrap_inner
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/threading.py", line 1002 in _bootstrap

Current thread 0x00007e3a3e73c740 (most recent call first):
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 456 in _conv_forward
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 460 in forward
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520 in _call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511 in _wrapped_call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/container.py", line 217 in forward
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520 in _call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511 in _wrapped_call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torchvision/models/optical_flow/raft.py", line 160 in forward
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520 in _call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511 in _wrapped_call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torchvision/models/optical_flow/raft.py", line 492 in forward
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520 in _call_impl
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511 in _wrapped_call_impl
File "/home/hadil/Work/flowmap/flowmap/flow/flow_predictor_raft.py", line 46 in forward
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/jaxtyping/_decorator.py", line 450 in wrapped_fn
File "/home/hadil/Work/flowmap/flowmap/flow/flow_predictor.py", line 87 in compute_bidirectional_flow
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/jaxtyping/_decorator.py", line 450 in wrapped_fn
File "/home/hadil/Work/flowmap/flowmap/flow/init.py", line 33 in compute_flows
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/jaxtyping/_decorator.py", line 450 in wrapped_fn
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115 in decorate_context
File "/home/hadil/Work/flowmap/flowmap/overfit.py", line 67 in overfit
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/core/utils.py", line 186 in run_job
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/_internal/hydra.py", line 119 in run
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/_internal/utils.py", line 458 in
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/_internal/utils.py", line 220 in run_and_report
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/_internal/utils.py", line 457 in _run_app
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/_internal/utils.py", line 394 in _run_hydra
File "/home/hadil/anaconda3/envs/flowmap/lib/python3.11/site-packages/hydra/main.py", line 94 in decorated_main
File "/home/hadil/Work/flowmap/flowmap/overfit.py", line 160 in
File "", line 88 in _run_code
File "", line 198 in _run_module_as_main

Extension modules: yaml._yaml, numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator, torch._C, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._nn, torch._C._sparse, torch._C._special, charset_normalizer.md, google._upb._message, psutil._psutil_linux, psutil._psutil_posix, matplotlib._c_internal_utils, PIL._imaging, matplotlib._path, kiwisolver._cext, matplotlib._image, PIL._imagingft, scipy._lib._ccallback_c, scipy.sparse._sparsetools, _csparsetools, scipy.sparse._csparsetools, scipy.linalg._fblas, scipy.linalg._flapack, scipy.linalg.cython_lapack, scipy.linalg._cythonized_array_utils, scipy.linalg._solve_toeplitz, scipy.linalg._flinalg, scipy.linalg._decomp_lu_cython, scipy.linalg._matfuncs_sqrtm_triu, scipy.linalg.cython_blas, scipy.linalg._matfuncs_expm, scipy.linalg._decomp_update, scipy.sparse.linalg._dsolve._superlu, scipy.sparse.linalg._eigen.arpack._arpack, scipy.sparse.csgraph._tools, scipy.sparse.csgraph._shortest_path, scipy.sparse.csgraph._traversal, scipy.sparse.csgraph._min_spanning_tree, scipy.sparse.csgraph._flow, scipy.sparse.csgraph._matching, scipy.sparse.csgraph._reordering, scipy.spatial._ckdtree, scipy._lib.messagestream, scipy.spatial._qhull, scipy.spatial._voronoi, scipy.spatial._distance_wrap, scipy.spatial._hausdorff, scipy.special._ufuncs_cxx, scipy.special._ufuncs, scipy.special._specfun, scipy.special._comb, scipy.special._ellip_harm_2, scipy.spatial.transform._rotation, PIL._imagingmath (total: 68)
[1] 27723 segmentation fault (core dumped) python -m flowmap.overfit dataset=images

[balabala932131]

There might not be any problem in the code. The error may occur due to your Python environment.
pip install torch==2.1.2+cu118 torchvision==0.16.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
conda install -c "nvidia/label/cuda-11.8.0" cuda-toolkit
its also work