Closed IanButterworth closed 10 months ago
@CarloLucibello could it be that the extensions aren't being loaded properly, even though gpu
is putting things on the gpu
Make sure you've moved model
to gpu
as well ;)
Oh ok. Is there a way to catch that and warn the user?
Also the docs to Metalhead don't appear to mention gpu
(even thought there are 3 results, but gpu
isn't found on any of the pages)
Warning the user is probably best handled at the NNlib level, that's what https://github.com/FluxML/NNlib.jl/issues/523 is tracking. We assume most users will have read through a page such as https://fluxml.ai/Flux.jl/stable/gpu/#Basic-GPU-Usage before using Metalhead (or any other Flux-based) models with GPUs. Those pages introduce the gpu
function and show how to apply it to both inputs and models.
Fair enough. I came straight to metalhead though and the flux models I've used before have internally managed gpu state, so it wasn't clear to me.
I'll try to come up with a docs PR
I'm not aware of any models which manage GPU state internally for parameters but not also for inputs. Let alone ones which then point people to functions like gpu
. Out of curiosity, do you have some examples? If this is happening in active downstream libraries, we could look into submitting some docs PRs.
ObjectDetector.jl is the one I had in mind. Note I'm a maintainer.. so I could be the issue!
The challenge with that is part of the post processing inside the model has been difficult to make runnable on the gpu.