allow NNlib v0.9 (#38)

* allow NNlib v0.9

* fix test

Author: CarloLucibello

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"
 

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