From 540b7366ec0edd711953223ef44bf342d691127f Mon Sep 17 00:00:00 2001
From: dsweber2 <david.weber2@gmail.com>
Date: Tue, 10 Sep 2019 00:54:49 -0700
Subject: [PATCH] make activations zygote friendly

---
 src/layers/basic.jl | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/layers/basic.jl b/src/layers/basic.jl
index 83eeee21..b4b869c5 100644
--- a/src/layers/basic.jl
+++ b/src/layers/basic.jl
@@ -45,17 +45,15 @@ end
 # it might be replaced in the future for better performance
 # see issue https://github.com/FluxML/Flux.jl/issues/702
 # Johnny Chen -- @johnnychen94
+# only slightly changed to better handle interaction with Zygote @dsweber2
 """
     activations(c::Chain, input)
 Calculate the forward results of each layers in Chain `c` with `input` as model input.
 """
 function activations(c::Chain, input)
-  rst = []
-  for l in c
-    x = get(rst, length(rst), input)
-    push!(rst, l(x))
-  end
-  return rst
+  buffed = accumulate!((x,y)->y(x), Zygote.Buffer([], length(c)),
+                       [l for l in c], dims=1, init=input) 
+  return copy(buffed)
 end