Correctly dispatch on type (#1)

sethaxen · May 14, 2020 · a635413 · a635413
1 parent 629fad5
commit a635413
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/buffer.jl b/src/buffer.jl
@@ -87,8 +87,8 @@ end
 grad_mut(b::Buffer) = fill!(similar(b.data, Any), nothing)
 grad_mut(b::Buffer{T}) where T<:Number = fill!(similar(b.data, float(T)), 0)
 
-function ZygoteRules._pullback(cx::AContext, ::typeof(Buffer), args...)
-  Buffer(args...), _ -> nothing
+function ZygoteRules._pullback(cx::AContext, ::Type{T}, args...) where {T<:Buffer}
+  T(args...), _ -> nothing
 end
 
 @adjoint function getindex(b::Buffer, i...)