google-research / maxim

[CVPR 2022 Oral] Official repository for "MAXIM: Multi-Axis MLP for Image Processing". SOTA for denoising, deblurring, deraining, dehazing, and enhancement.
https://arxiv.org/abs/2201.02973
Apache License 2.0
977 stars 105 forks source link

No GPU found #21

Closed hs0817 closed 1 year ago

hs0817 commented 1 year ago

Hello, I installed the environment and projects according to requirements. txt, but there was no GPU in the test process. My CUDA version was 11.2, and the GPU was RTX3090. Can you help me explain this situation? 1659858182299

Walterfdh commented 1 year ago

You should use the jaxlib's cuda version.

hs0817 commented 1 year ago

Thank you for your quick response.