diff --git a/Project.toml b/Project.toml
index ebbe4ae..aaef604 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "1.0.0"
+version = "1.0.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/src/autodiff.jl b/src/autodiff.jl
index bdc4d2a..1dc41f0 100644
--- a/src/autodiff.jl
+++ b/src/autodiff.jl
@@ -108,7 +108,7 @@ julia> test_gradients(f, 1.0, x, nothing)
 ```
 """
 function test_gradients(f, args...; skip_backends=[], broken_backends=[], kwargs...)
-    on_gpu = get_device_type(args) isa AbstractGPUDevice
+    on_gpu = get_device_type(args) <: AbstractGPUDevice
     total_length = mapreduce(__length, +, Functors.fleaves(args); init=0)
 
     # Choose the backends to test