Scalar indexing triggered by remove_edges

[alonsoC1s]

The function remove_edges triggers a scalar indexing exception when the graph resides on the GPU. Other similar functions, like add_edges, don't have this problem

A MWE using CUDA:

using CUDA, GraphNeuralNetworks, SparseArrays, Flux
G = GNNGraph(A, ndata=randn(20)) |> Flux.gpu

julia> remove_edges(G, cu([3, 4, 5]))
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

GraphNeuralNetworks version is v1.1.0.

Could it be that this method is specialized on GNNGraph{<:COO_t}? From a quick look at the function's source, it doesn't look like it does anything GPU-illegal

The culprit seems to be:

mask_to_keep = trues(length(s))
mask_to_keep[edges_to_remove] .= false # This line

This broadcast is not an issue when mask_to_keep is a GPU vector as well

[CarloLucibello]

so the fix seems easy, initialize mask_to_keep on gpu. Want to make a PR?