Differentiating through my custom struct its restructuring throws an error

[JTaets]

Title says all.

I could boil down my errors to this mwe.

using Flux

struct Model
    model
    offset
end
@Flux.functor Model (model,)
(m::Model)(x) =m.model(x+m.offset)

model = Model(Dense(1,1),[4])
p,re = Flux.destructure(model)
println(length(p)) #2 parameters
Flux.gradient((p)->sum(re(p)([1])),p) #doesn't work

The last lines throws an error DimensionMismatch("variable with size(x) == (2,) cannot have a gradient with size(dx) == (3,)")

[DhairyaLGandhi]

Could you try with @functor Model

[JTaets]

Then it works, but I don't want this offset to be differentiable

[ToucheSir]

Try (m::Model)(x) = m.model(x + Zygote.dropgrad(m.offset))¹ with your original @functor Model (model,). Just tested and it works for me locally.

Footnotes

1. https://fluxml.ai/Zygote.jl/latest/utils/#Zygote.dropgrad ↩

[JTaets]

This also works for my actual case. Thank you!