Cannot save Flux gpu model weights

[yaomz16]

Hi, I wonder whether @save can save a gpu model weights properly. Here's a MWE (MWE.jl):

`using DifferentialEquations, Flux, DiffEqFlux, NPZ, Printf, CUDA
using BSON: @save, @load
CUDA.allowscalar(false)

dudt = Chain(Conv((7,7),1=>1, swish, pad=3)) |> gpu
n_ode = NeuralODE(dudt, (0.0f0, 0.1f0),Tsit5(),saveat=[0.0f0, 0.1f0],reltol=1e-7,abstol=1e-9) |> gpu
weights = Flux.params(n_ode)
@save "n_ode_weights_tmp.bson" weights
print("Try to load the model\n")
print(CUDA.has_cuda_gpu(),"\n")
@load "n_ode_weights_tmp.bson" weights
Flux.loadparams!(n_ode, weights)
`

When I run this, I will get the following error:
Try to load the model
true
ERROR: LoadError: ArgumentError: copying between arrays from different contexts
Stacktrace:
[1] unsafe_copyto!(dest::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, doffs::Int64, src::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, soffs::Int64, n::Int64)
@ CUDA ~/.julia/packages/CUDA/YpW0k/src/array.jl:424
[2] copyto!
@ ~/.julia/packages/CUDA/YpW0k/src/array.jl:373 [inlined]
[3] copyto!(dest::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, src::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer})
@ CUDA ~/.julia/packages/CUDA/YpW0k/src/array.jl:377
[4] loadparams!(m::Function, xs::Params)
@ Flux ~/.julia/packages/Flux/ZnXxS/src/functor.jl:61
[5] top-level scope
@ /ocean/projects/cts180021p/mingzeya/julia_test/MWE.jl:13
in expression starting at /ocean/projects/cts180021p/mingzeya/julia_test/MWE.jl:13

versioninfo:
julia> versioninfo()
Julia Version 1.6.3
Commit ae8452a9e0 (2021-09-23 17:34 UTC)
Platform Info:
OS: Linux (x86_64-pc-linux-gnu)
CPU: Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-11.0.1 (ORCJIT, cascadelake)

pkg status:
(@v1.6) pkg> st
Status ~/.julia/environments/v1.6/Project.toml
[fbb218c0] BSON v0.3.4
[052768ef] CUDA v3.5.0
[aae7a2af] DiffEqFlux v1.43.0
[0c46a032] DifferentialEquations v6.19.0
[587475ba] Flux v0.12.7
[15e1cf62] NPZ v0.4.1
[e88e6eb3] Zygote v0.6.26
[de0858da] Printf

[DhairyaLGandhi]

You need to move the data over to cpu first since CUDA really only stores the pointers, not the arrays themselves.

…

On Tue, Oct 19, 2021, 00:03 Archie Yao ***@***.***> wrote: Hi, I wonder whether @save <https://github.com/save> can save a gpu model weights properly. Here's a MWE (MWE.jl): `using DifferentialEquations, Flux, DiffEqFlux, NPZ, Printf, CUDA using BSON: @save <https://github.com/save>, @load <https://github.com/load> CUDA.allowscalar(false) dudt = Chain(Conv((7,7),1=>1, swish, pad=3)) |> gpu n_ode = NeuralODE(dudt, (0.0f0, 0.1f0),Tsit5(),saveat=[0.0f0, 0.1f0],reltol=1e-7,abstol=1e-9) |> gpu weights = Flux.params(n_ode) @save <https://github.com/save> "n_ode_weights_tmp.bson" weights print("Try to load the model\n") print(CUDA.has_cuda_gpu(),"\n") @load <https://github.com/load> "n_ode_weights_tmp.bson" weights Flux.loadparams!(n_ode, weights) ` When I run this, I will get the following error: Try to load the model true ERROR: LoadError: ArgumentError: copying between arrays from different contexts Stacktrace: [1] unsafe_copyto!(dest::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, doffs::Int64, src::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, soffs::Int64, n::Int64) @ CUDA ~/.julia/packages/CUDA/YpW0k/src/array.jl:424 [2] copyto! @ ~/.julia/packages/CUDA/YpW0k/src/array.jl:373 [inlined] [3] copyto!(dest::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, src::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}) @ CUDA ~/.julia/packages/CUDA/YpW0k/src/array.jl:377 [4] loadparams!(m::Function, xs::Params) @ Flux ~/.julia/packages/Flux/ZnXxS/src/functor.jl:61 [5] top-level scope @ /ocean/projects/cts180021p/mingzeya/julia_test/MWE.jl:13 in expression starting at /ocean/projects/cts180021p/mingzeya/julia_test/MWE.jl:13 versioninfo: julia> versioninfo() Julia Version 1.6.3 Commit ae8452a9e0 (2021-09-23 17:34 UTC) Platform Info: OS: Linux (x86_64-pc-linux-gnu) CPU: Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz WORD_SIZE: 64 LIBM: libopenlibm LLVM: libLLVM-11.0.1 (ORCJIT, cascadelake) pkg status: — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#104>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AJOZVVIEPOPWVKL2TQGZQATUHRR6HANCNFSM5GHJQESQ> . Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.