LoadError: indexing not defined for FluxJS.StagedArray{Float64,2}

I tried to run this basic code for a cppn
```julia
using Flux
using Images
using FluxJS
function transform(z)
x_dim = 256
y_dim = 256
n_points = x_dim*y_dim
scale = 8
z_dim = 32
net_size = 32
x_range = scale*(collect(0:(x_dim-1))-(x_dim-1)/2.0)/(x_dim-1)/0.5
y_range = scale*(collect(0:(y_dim-1))-(y_dim-1)/2.0)/(y_dim-1)/0.5
x_mat = ones(y_dim, 1)*reshape(x_range, 1, x_dim)
y_mat = reshape(y_range, y_dim, 1)*ones((1, x_dim))
r_mat = sqrt.(x_mat.^2 + y_mat.^2)
x_mat = reshape(transpose(x_mat), n_points)
y_mat = reshape(transpose(y_mat), n_points)
r_mat = reshape(transpose(r_mat), n_points)

unit = Dense(n_points, net_size)
unit(x_mat) + unit(y_mat) + unit(r_mat)
end
model = Chain(
transform,
Dense(32, 32, tanh),
Dense(32, 32, tanh),
Dense(32, 32, tanh),
Dense(32, 32, σ)
)
@code_js model(rand(32))
```