fused-local-corr/local_corr/ops.py has VRAM leak when used with RoMaV2

[ferrouswheel]

When running 100s of images through RoMA installed with romav2[fused-local-corr] the VRAM continues to creep upwards.

The solution appears to be changing these lines in fused-local-corr/local_corr/ops.py

89      saved_a = a if ctx.needs_input_grad[0] else None                                                                                                
90 -    saved_b = b if ctx.needs_input_grad[0] else None                                                                                                
91 -    saved_c = c if ctx.needs_input_grad[0] else None                                                                                                
90 +    saved_b = b if ctx.needs_input_grad[1] else None                                                                                                
91 +    saved_c = c if ctx.needs_input_grad[2] else None

Since fused-local-corr is on a gitlab instance I have no access to, I am unable to submit a patch directly to https://gitlab.liu.se/johed13/fused-local-corr

[ferrouswheel]

Possibly related to #17

[LEFTazs]

I also have the same OOM issue, and this patch didn't solve it for me.

I have 30 frames and by the time I get to the second frame, my 12gb VRAM fills up. It works with 3 frames only. It is visible with the pytorch memory_viz, that fused-local-corr has variables that are not being freed up between iterations.

[ferrouswheel]

@LEFTazs I recall that I found it very finnicky to ensure that uv was using the correctly patched version after already installing it.

You could try creating a new .venv and ensure that fuse_local_corr is pointing at your local copy in the pyproject.toml ?

May have also been a need to override the uv cache, but alas I can't remember. That patch did solve it for me, but I spent a fair amount of time JUST trying to get uv to actually install it over the previous version 😓