use Flux.@layer instead of Flux.@functor

This would get us fancy printing with model parameters count and suppress the warnings 
```
┌ Warning: This type should probably now use `Flux.@layer` instead of `@functor`: WithGraph{SAGEConv{Matrix{Float32}, Vector{Float32}, typeof(identity), typeof(mean)}, GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}
```