fix: skip FiniteDiff on GPU too slow

LuxDL · Jul 28, 2024 · fd8db73 · fd8db73 · avik-pal · Jul 28, 2024
1 parent e8e62aa
commit fd8db73
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/autodiff.jl b/src/autodiff.jl
@@ -119,13 +119,14 @@ function test_gradients(f, args...; skip_backends=[], broken_backends=[], kwargs
         if AutoForwardDiff() ∉ skip_backends && total_length ≤ 100
             push!(backends, AutoForwardDiff())
         end
+        if AutoFiniteDiff() ∉ skip_backends && total_length ≤ 100
+            push!(backends, AutoFiniteDiff())
+        end
+        # TODO: Move Enzyme out of here once it supports GPUs
         if AutoEnzyme() ∉ skip_backends && ENZYME_TESTING_ENABLED
             push!(backends, AutoEnzyme())
         end
     end
-    if AutoFiniteDiff() ∉ skip_backends && total_length ≤ 100
-        push!(backends, AutoFiniteDiff())
-    end
     AutoTracker() ∉ skip_backends && push!(backends, AutoTracker())
 
     # Test the gradients