Skip to content

Commit

Permalink
fix: remove deprecated API
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Sep 21, 2024
1 parent c8e8c91 commit 2255638
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions test/gpu/mixed_gpu_cpu_adjoint.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
using SciMLSensitivity, OrdinaryDiffEq
using Lux, LuxCUDA, Test, Zygote, Random, LinearAlgebra, ComponentArrays

const gdev = gpu_device()

CUDA.allowscalar(false)

H = CuArray(rand(Float32, 2, 2))
Expand Down Expand Up @@ -42,10 +44,10 @@ grad = Zygote.gradient(cost, p)[1]
rng = MersenneTwister(1234)
m = 32
n = 16
Z = randn(rng, Float32, (n, m)) |> gpu
Z = randn(rng, Float32, (n, m)) |> gdev
𝒯 = 2.0f0
Δτ = 1.0f-1
ca_init = [zeros(1); ones(m)] |> gpu
ca_init = [zeros(1); ones(m)] |> gdev

function f(ca, Z, t)
a = ca[2:end]
Expand All @@ -54,7 +56,7 @@ function f(ca, Z, t)
Ka_unit = Z' * w_unit
z_unit = dot(abs.(Ka_unit), a_unit)
aKa_over_z = a .* Ka_unit / z_unit
[sum(aKa_over_z) / m; -abs.(aKa_over_z)] |> gpu
[sum(aKa_over_z) / m; -abs.(aKa_over_z)] |> gdev
end

function c(Z)
Expand Down

0 comments on commit 2255638

Please sign in to comment.