feat: add option for mixed precision dtype to enable running on cpus

Parskatt / RoMa

[CVPR 2024] RoMa: Robust Dense Feature Matching; RoMa is the robust dense feature matcher capable of estimating pixel-dense warps and reliable certainties for almost any image pair.

https://parskatt.github.io/RoMa/

MIT License

556 stars 43 forks source link

feat: add option for mixed precision dtype to enable running on cpus #31

Closed dgcnz closed 5 months ago

dgcnz commented 5 months ago

Solves https://github.com/Parskatt/RoMa/issues/26

Changes:

Add optional parameter amp_dtype

One can now run RoMa on cpu by specifying amp_dtype=torch.float32 like so:

roma_model = roma_outdoor(device=device, amp_dtype=torch.float32)

Parskatt commented 5 months ago

Seems good, did you check that it works? Relatedly I should probably add some tests haha.

Parskatt commented 5 months ago

I'll merge it anyway, shouldn't be any issues I think.