Running model multiple times gets much slower?

[FR13ndSDP]

Here is a simple piece of code:

using JLD2, Lux, LuxCUDA

@load "./NN/luxmodel.jld2" model ps st

ps = ps |> gpu_device()

input = CUDA.ones(Float64, 9, 1024*256)

@time for i = 1:10
    Lux.apply(model, input, ps, st)
end

I just found that if the model is being run only once, the time it takes is 0.000380 seconds (4.81 k allocations: 218.922 KiB)

But if it is being run 10 times inside a for loop, the time it takes is 0.380598 seconds (8.69 k allocations: 404.797 KiB, 0.60% gc time).

I think this problem might be related to GC, how can I avoid it?

[FR13ndSDP]

Just found out that I should use CUDA.@time instead of @time. Converting input to Float32 and running the model using Lux.apply(model, cu(input), ps, st) is much faster.