allow NNlib v0.9 (#38)

* allow NNlib v0.9 * fix test
FluxML · Jun 16, 2023 · a531d55 · a531d55
1 parent 469b192
commit a531d55
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -12,11 +12,11 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 
 [compat]
 Adapt = "3.0"
-CUDA = "3.8"
+CUDA = "4"
 ChainRulesCore = "1.13"
 Compat = "4.2"
 GPUArraysCore = "0.1.0"
-NNlib = "0.8"
+NNlib = "0.8, 0.9"
 Zygote = "0.6.35"
 julia = "1.6"
 

diff --git a/test/gpu.jl b/test/gpu.jl
@@ -23,7 +23,11 @@ end
   @test (repr("text/plain", y); true)
 
   gA = rand(3, 2) |> cu;
-  @test_broken gradient(A -> sum(A * y), gA)[1] isa CuArray  # fails with JLArray, bug in Zygote?
+  if VERSION >= v"1.9" && CUDA.functional()
+    @test gradient(A -> sum(A * y), gA)[1] isa CuArray 
+  else
+    @test_broken gradient(A -> sum(A * y), gA)[1] isa CuArray  # fails with JLArray, bug in Zygote?
+  end
 end
 
 @testset "onehotbatch(::CuArray, ::UnitRange)" begin