Good day all,
I am struggling a bit with something simple, I hope. I would like to use
Zygote.pushforward to AD neural network outputs wrt inputs. For standard fully connected nets everything works fine, but as soon as I use a modified network with residual connections I get the following error message:
Generic broadcast of * not support yet
Is there anyway to get around this? Below is a MWE. Any help will be appreciated.
using Zygote using Statistics using Flux data = randn(2,10) data_ones = ones(2,10) fc1 = Dense(2,50) fc2 = Dense(50,50) fc3 = Dense(50,6) fc_U = Dense(2,50) function model(x) U = fc_U(x) h = fc1(x) h = fc2(h) .* U h = fc3(h) h end test = model(data) grad_fun = Zygote.pushforward(model, data) grad_fun(data_ones)