You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
input_shape =rand(28,28,1,1);
body =Chain(
StandardConv((3,3),1=>1),
x ->reshape(x,:,size(x,4)),
)
body_shape =size(body(input_shape),1)
Flux.reset!(body)
head =Chain(
Dense(body_shape,512),
Dense(512,10)
)
model =Chain(
body,
head
);
Flux.reset!(model);
The loss
functionloss(x,y)
y_pred =model(x) |> softmax
Flux.logitcrossentropy(y_pred,y)
end#dummy input
x =rand(28,28,1,10)
y = Flux.onehotbatch(collect(1:10),1:10)
θ =params(model);
grads = Tracker.gradient(() ->loss(x,y) ,θ)
If you comment out layer.norm(x) the gradients work fine?
The text was updated successfully, but these errors were encountered:
When I try to stick BatchNorm in a custom layer it causes an error when calculating the gradient.
The layer looks like this
The model
The loss
If you comment out
layer.norm(x)
the gradients work fine?The text was updated successfully, but these errors were encountered: