Set model parameters as Dictionary and define some as trainable

General Usage
1

I defined the following model with a dictionary containing weights

mutable struct Affine
  params
end

Affine(in::Integer, out::Integer) =
  Affine(Dict("W"=>randn(out, in), ("b"=>randn(out))))

# Overload call, so the object can be used as a function
(m::Affine)(x) = m.params["W"] * x .+ m.params["b"]

a = Affine(1, 1)

#a.params
#Dict{String,Array{Float64,N} where N} with 2 entries:
#  "W" => [0.559984]
#  "b" => [0.841258]

Then I tried to specify params[“W”] as trainable

Flux.@functor Affine
Flux.trainable(a::Affine) = (a.params["W"],)

grads = Flux.gradient(() -> a(3)[1], params(a))
for p in grads.grads
    println(p)
    println("--")
end

# Output
Pair{Any,Any}([0.5599836013031454], nothing)
--
Pair{Any,Any}(Dict{String,Array{Float64,N} where N}("W" => [0.5599836013031454],"b" => [0.841258238965482]), Dict{Any,Any}("W" => [3.0],"b" => [1.0]))
--

But gradient seems to be taken with respect to both parameters.
Can I choose only one I need?