# Lux.jl with GPU error

**URL:** <https://discourse.julialang.org/t/lux-jl-with-gpu-error/121060>\
**Category:** Machine Learning\
**Created:** [October 8, 2024, 12:06pm UTC](https://discourse.julialang.org/t/lux-jl-with-gpu-error/121060 "2024-10-08T12:06:20Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Sushrut\_Deshpande](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sushrut_deshpande/32/52754_2.png) [@Sushrut\_Deshpande](https://discourse.julialang.org/u/Sushrut_Deshpande)\
**Post date:** [October 8, 2024, 12:06pm UTC](https://discourse.julialang.org/t/lux-jl-with-gpu-error/121060/1 "2024-10-08T12:06:20Z")

</div>

Hello,  
I have a toy problem to fit a NN for `sin(x)` using the GPU for Lux.jl - LuxCUDA.jl

Here is the code I am running (very similar to the documentation):

```julia
using Lux, Optimisers,Zygote, Random
using LuxCUDA

x = rand(50);
y = sin.(x)

rng = MersenneTwister()
model = Chain(Dense(1=>20,tanh),Dense(20=>40),Dense(40=>40),Dense(40=>1,tanh))
opt = Adam()
const loss_function = MSELoss()
const dev_gpu = gpu_device()
const dev_cpu = cpu_device()

ps, st = Lux.setup(rng, model) .|> dev_gpu

tstate = Training.TrainState(model, ps, st, opt)
vjp_rule = AutoZygote()

using ProgressMeter

function main(tstate::Training.TrainState, vjp, data, epochs)
    data = data .|> dev_gpu
    progressmeter = Progress(epochs; showspeed=true)
    for epoch in 1:epochs
        _, loss, _, tstate = Training.single_train_step!(vjp, loss_function, data, tstate)
        next!(progressmeter; showvalues = [(:Epoch,epoch), (:Loss,loss)])
    end
    return tstate
end

@info "Training"
tstate = main(tstate, vjp_rule, (x', y'), 10000)

```

This code works fine for `cpu_device()` but for `gpu_device()` I get the following error:

```julia
ERROR: LoadError: MethodError: no method matching cublasLt_fused_dense!(::CuArray{Float32, 2, CUDA.DeviceMemory}, ::typeof(tanh_fast), ::CuArray{Float32, 2, CUDA.DeviceMemory}, ::LinearAlgebra.Adjoint{Float32, CuArray{…}}, ::CuArray{Float32, 1, CUDA.DeviceMemory})

```

I am not sure what is my error.

Thank you

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 8, 2024, 12:22pm UTC](https://discourse.julialang.org/t/lux-jl-with-gpu-error/121060/2 "2024-10-08T12:22:09Z")

</div>

> [@Sushrut\_Deshpande](#):
>
> `tstate = main(tstate, vjp_rule, (x', y'), 10000)`

Try changing that line to `tstate = main(tstate, vjp_rule, (reshape(x, 1, 50), reshape(y, 1, 50)), 10000)`

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 10, 2024, 7:11pm UTC](https://discourse.julialang.org/t/lux-jl-with-gpu-error/121060/3 "2024-10-10T19:11:56Z")

</div>

[fix: relax cublaslt types by avik-pal · Pull Request #173 · LuxDL/LuxLib.jl · GitHub](https://github.com/LuxDL/LuxLib.jl/pull/173) will fix this for `x'` and `y'`.

---

<div class="post-metadata">

**Author:** ![Sushrut\_Deshpande](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sushrut_deshpande/32/52754_2.png) [@Sushrut\_Deshpande](https://discourse.julialang.org/u/Sushrut_Deshpande)\
**Post date:** [October 10, 2024, 7:25pm UTC](https://discourse.julialang.org/t/lux-jl-with-gpu-error/121060/4 "2024-10-10T19:25:53Z")

</div>

Thanks!!