Add workaround for <FluxML/Zygote.jl#962>

marcoq · Jun 9, 2021 · 5771aa3 · 5771aa3
1 parent 24bc905
commit 5771aa3
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/NeuralSurrogate.jl b/src/NeuralSurrogate.jl
@@ -1,5 +1,11 @@
 using Flux
 using Flux: @epochs
+
+### NOTICE: the next thee lines are a workaround for <https://github.com/FluxML/Zygote.jl/pull/962>
+using Base
+using Zygote
+Base.setindex!(dict::IdDict, dx::Zygote.OneElement, x) = dict[x] = collect(dx)
+
 mutable struct NeuralSurrogate{X,Y,M,L,O,P,N,A,U} <: AbstractSurrogate
     x::X
     y::Y