diff --git a/GNNGraphs/docs/make.jl b/GNNGraphs/docs/make.jl
index f0dc2ea44..445fdbb4c 100644
--- a/GNNGraphs/docs/make.jl
+++ b/GNNGraphs/docs/make.jl
@@ -19,7 +19,7 @@ mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/ma
 
 makedocs(;
     modules = [GNNGraphs],
-    doctest = false, # TODO enable doctest
+    doctest = true, # TODO enable doctest
     format = Documenter.HTML(; mathengine, 
                     prettyurls = get(ENV, "CI", nothing) == "true", 
                     assets = [],
@@ -30,9 +30,9 @@ makedocs(;
     
     "Guides" => [
         "Graphs" => [
-        "guides/gnngraph.md", 
-        "guides/heterograph.md", 
-        "guides/temporalgraph.md"
+            "guides/gnngraph.md", 
+            "guides/heterograph.md", 
+            "guides/temporalgraph.md"
         ],
         "Datasets" => "guides/datasets.md",
     ],
diff --git a/GNNGraphs/docs/src/guides/gnngraph.md b/GNNGraphs/docs/src/guides/gnngraph.md
index c62e279fa..bea1c28ce 100644
--- a/GNNGraphs/docs/src/guides/gnngraph.md
+++ b/GNNGraphs/docs/src/guides/gnngraph.md
@@ -155,11 +155,10 @@ that contains the total number of the original nodes
 and where the original graphs are disjoint subgraphs.
 
 ```julia
-using Flux
-using Flux: DataLoader
+using MLUtils
 
 data = [rand_graph(10, 30, ndata=rand(Float32, 3, 10)) for _ in 1:160]
-gall = Flux.batch(data)
+gall = MLUtils.batch(data)
 
 # gall is a GNNGraph containing many graphs
 @assert gall.num_graphs == 160 
@@ -172,7 +171,7 @@ g23 = getgraph(gall, 2:3)
 @assert g23.num_nodes == 20   # 10 nodes x 2 graphs
 @assert g23.num_edges == 60  # 30 undirected edges X 2 graphs
 
-# We can pass a GNNGraph to Flux's DataLoader
+# We can pass a GNNGraph to MLUtils' DataLoader
 train_loader = DataLoader(gall, batchsize=16, shuffle=true)
 
 for g in train_loader
@@ -193,7 +192,7 @@ an option for mini-batch iteration, the recommended way for better performance i
 to pass an array of graphs directly and set the `collate` option to `true`:
 
 ```julia
-using Flux: DataLoader
+using MLUtils: DataLoader
 
 data = [rand_graph(10, 30, ndata=rand(Float32, 3, 10)) for _ in 1:320]
 
@@ -220,7 +219,7 @@ g′ = add_edges(g, [1, 2], [2, 3]) # add edges 1->2 and 2->3
 Move a `GNNGraph` to a CUDA device using `Flux.gpu` method. 
 
 ```julia
-using CUDA, Flux
+using Flux, CUDA # or using Metal or using AMDGPU 
 
 g_gpu = g |> Flux.gpu
 ```
diff --git a/GNNGraphs/docs/src/guides/heterograph.md b/GNNGraphs/docs/src/guides/heterograph.md
index 2347b5844..7999e2a48 100644
--- a/GNNGraphs/docs/src/guides/heterograph.md
+++ b/GNNGraphs/docs/src/guides/heterograph.md
@@ -13,7 +13,9 @@ the type [`GNNHeteroGraph`](@ref).
 ## Creating a Heterograph
 
 A heterograph can be created empty or by passing pairs `edge_type => data` to the constructor.
-```jldoctest
+```jldoctest hetero
+julia> using GNNGraphs
+
 julia> g = GNNHeteroGraph()
 GNNHeteroGraph:
   num_nodes: Dict()
@@ -31,7 +33,7 @@ GNNHeteroGraph:
   num_edges: Dict((:user, :rate, :movie) => 4)
 ```
 New relations, possibly with new node types, can be added with the function [`add_edges`](@ref).
-```jldoctest
+```jldoctest hetero
 julia> g = add_edges(g, (:user, :like, :actor) => ([1,2,3,3,3], [3,5,1,9,4]))
 GNNHeteroGraph:
   num_nodes: Dict(:actor => 9, :movie => 13, :user => 3)
@@ -40,7 +42,7 @@ GNNHeteroGraph:
 See [`rand_heterograph`](@ref), [`rand_bipartite_heterograph`](@ref)
 for generating random heterographs. 
 
-```jldoctest
+```jldoctest hetero
 julia> g = rand_bipartite_heterograph((10, 15), 20)
 GNNHeteroGraph:
   num_nodes: Dict(:A => 10, :B => 15)
@@ -50,7 +52,7 @@ GNNHeteroGraph:
 ## Basic Queries
 
 Basic queries are similar to those for homogeneous graphs:
-```jldoctest
+```jldoctest hetero
 julia> g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
 GNNHeteroGraph:
   num_nodes: Dict(:movie => 13, :user => 3)
@@ -84,7 +86,7 @@ julia> g.etypes
 ## Data Features
 
 Node, edge, and graph features can be added at construction time or later using:
-```jldoctest
+```jldoctest hetero
 # equivalent to g.ndata[:user][:x] = ...
 julia> g[:user].x = rand(Float32, 64, 3);
 
@@ -106,10 +108,10 @@ GNNHeteroGraph:
 
 ## Batching
 Similarly to graphs, also heterographs can be batched together.
-```jldoctest
+```jldoctest hetero
 julia> gs = [rand_bipartite_heterograph((5, 10), 20) for _ in 1:32];
 
-julia> Flux.batch(gs)
+julia> MLUtils.batch(gs)
 GNNHeteroGraph:
   num_nodes: Dict(:A => 160, :B => 320)
   num_edges: Dict((:A, :to, :B) => 640, (:B, :to, :A) => 640)
@@ -118,7 +120,7 @@ GNNHeteroGraph:
 Batching is automatically performed by the [`DataLoader`](https://fluxml.ai/Flux.jl/stable/data/mlutils/#MLUtils.DataLoader) iterator
 when the `collate` option is set to `true`.
 
-```jldoctest
+```jldoctest hetero
 using Flux: DataLoader
 
 data = [rand_bipartite_heterograph((5, 10), 20, 
diff --git a/GNNGraphs/docs/src/guides/temporalgraph.md b/GNNGraphs/docs/src/guides/temporalgraph.md
index 560cfa8d6..40d817cdf 100644
--- a/GNNGraphs/docs/src/guides/temporalgraph.md
+++ b/GNNGraphs/docs/src/guides/temporalgraph.md
@@ -6,7 +6,9 @@ Temporal Graphs are graphs with time varying topologies and  features. In GNNGra
 
 A temporal graph can be created by passing a list of snapshots to the constructor. Each snapshot is a [`GNNGraph`](@ref). 
 
-```jldoctest
+```jldoctest temporal
+julia> using GNNGraphs
+
 julia> snapshots = [rand_graph(10,20) for i in 1:5];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -18,14 +20,14 @@ TemporalSnapshotsGNNGraph:
 
 A new temporal graph can be created by adding or removing snapshots to an existing temporal graph. 
 
-```jldoctest
+```jldoctest temporal
 julia> new_tg = add_snapshot(tg, 3, rand_graph(10, 16)) # add a new snapshot at time 3
 TemporalSnapshotsGNNGraph:
   num_nodes: [10, 10, 10, 10, 10, 10]
   num_edges: [20, 20, 16, 20, 20, 20]
   num_snapshots: 6
 ```
-```jldoctest
+```jldoctest temporal
 julia> snapshots = [rand_graph(10,20), rand_graph(10,14), rand_graph(10,22)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -43,7 +45,7 @@ TemporalSnapshotsGNNGraph:
 
 See [`rand_temporal_radius_graph`](@ref) and [`rand_temporal_hyperbolic_graph`](@ref) for generating random temporal graphs. 
 
-```jldoctest
+```jldoctest temporal
 julia> tg = rand_temporal_radius_graph(10, 3, 0.1, 0.5)
 TemporalSnapshotsGNNGraph:
   num_nodes: [10, 10, 10]
@@ -54,7 +56,7 @@ TemporalSnapshotsGNNGraph:
 ## Basic Queries
 
 Basic queries are similar to those for [`GNNGraph`](@ref)s:
-```jldoctest
+```jldoctest temporal
 julia> snapshots = [rand_graph(10,20), rand_graph(10,14), rand_graph(10,22)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -94,7 +96,7 @@ GNNGraph:
 A temporal graph can store global feature for the entire time series in the `tgdata` filed.
 Also, each snapshot can store node, edge, and graph features in the `ndata`, `edata`, and `gdata` fields, respectively. 
 
-```jldoctest
+```jldoctest temporal
 julia> snapshots = [rand_graph(10,20; ndata = rand(3,10)), rand_graph(10,14; ndata = rand(4,10)), rand_graph(10,22; ndata = rand(5,10))]; # node features at construction time
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots);
@@ -124,22 +126,4 @@ julia> [g.x for g in tg.snapshots]; # same vector as above, now accessing
                                    # the x feature directly from the snapshots
 ```
 
-## Graph convolutions on TemporalSnapshotsGNNGraph
-
-A graph convolutional layer can be applied to each snapshot independently, in the next example we apply a `GINConv` layer to each snapshot of a `TemporalSnapshotsGNNGraph`.  
-
-```jldoctest
-julia> using GNNGraphs, Flux
-
-julia> snapshots = [rand_graph(10, 20; ndata = rand(3, 10)), rand_graph(10, 14; ndata = rand(3, 10))];
-
-julia> tg = TemporalSnapshotsGNNGraph(snapshots);
-
-julia> m = GINConv(Dense(3 => 1), 0.4);
-
-julia> output = m(tg, tg.ndata.x);
-
-julia> size(output[1])
-(1, 10)
-```
 
diff --git a/GNNGraphs/src/datastore.jl b/GNNGraphs/src/datastore.jl
index 9581e266d..1bfa37e00 100644
--- a/GNNGraphs/src/datastore.jl
+++ b/GNNGraphs/src/datastore.jl
@@ -37,7 +37,7 @@ DataStore() with 2 elements:
 The `DataStore` has an interface similar to both dictionaries and named tuples.
 Arrays can be accessed and added using either the indexing or the property syntax:
 
-```jldoctest
+```jldoctest datastore
 julia> ds = DataStore(x = ones(Float32, 2, 3), y = zeros(Float32, 3))
 DataStore() with 2 elements:
   y = 3-element Vector{Float32}
@@ -49,20 +49,21 @@ julia> ds.x   # same as `ds[:x]`
  1.0  1.0  1.0
 
 julia> ds.z = zeros(Float32, 3)  # Add new feature array `z`. Same as `ds[:z] = rand(Float32, 3)`
-3-element Vector{Float64}:
-0.0
-0.0
-0.0
+3-element Vector{Float32}:
+ 0.0
+ 0.0
+ 0.0
 ```
 
 The `DataStore` can be iterated over, and the keys and values can be accessed
 using `keys(ds)` and `values(ds)`. `map(f, ds)` applies the function `f`
 to each feature array:
 
-```jldoctest
-julia> ds = DataStore(a = zeros(2), b = zeros(2));
-
+```jldoctest datastore
 julia> ds2 = map(x -> x .+ 1, ds)
+DataStore() with 2 elements:
+  a = 2-element Vector{Float64}
+  b = 2-element Vector{Float64}
 
 julia> ds2.a
 2-element Vector{Float64}:
diff --git a/GNNGraphs/src/gnngraph.jl b/GNNGraphs/src/gnngraph.jl
index a9af576e2..6ebdee8f7 100644
--- a/GNNGraphs/src/gnngraph.jl
+++ b/GNNGraphs/src/gnngraph.jl
@@ -66,7 +66,7 @@ functionality from that library.
 # Examples
 
 ```julia
-using GraphNeuralNetworks
+using GNNGraphs
 
 # Construct from adjacency list representation
 data = [[2,3], [1,4,5], [1], [2,5], [2,4]]
@@ -101,11 +101,8 @@ g.edata.e # or just g.e
 # Both source and target are vectors of length num_edges
 source, target = edge_index(g)
 ```
-A `GNNGraph` can be sent to the GPU using e.g. Flux's `gpu` function:
-```
-# Send to gpu
-using Flux, CUDA
-g = g |> Flux.gpu
+A `GNNGraph` can be sent to the GPU, for example by using Flux.jl's `gpu` function
+or MLDataDevices.jl's utilities. 
 ```
 """
 struct GNNGraph{T <: Union{COO_T, ADJMAT_T}} <: AbstractGNNGraph{T}
diff --git a/GNNGraphs/src/gnnheterograph.jl b/GNNGraphs/src/gnnheterograph.jl
index 72d67b34b..7e55e76e7 100644
--- a/GNNGraphs/src/gnnheterograph.jl
+++ b/GNNGraphs/src/gnnheterograph.jl
@@ -37,7 +37,7 @@ It is similar to [`GNNGraph`](@ref) but nodes and edges are of different types.
 # Examples
 
 ```julia
-julia> using GraphNeuralNetworks
+julia> using GNNGraphs
 
 julia> nA, nB = 10, 20;
 
diff --git a/GNNGraphs/src/mldatasets.jl b/GNNGraphs/src/mldatasets.jl
index 1f2bf7139..df072ca9a 100644
--- a/GNNGraphs/src/mldatasets.jl
+++ b/GNNGraphs/src/mldatasets.jl
@@ -8,7 +8,7 @@ Convert a graph dataset from the package MLDatasets.jl into one or many [`GNNGra
 # Examples
 
 ```jldoctest
-julia> using MLDatasets, GraphNeuralNetworks
+julia> using MLDatasets, GNNGraphs
 
 julia> mldataset2gnngraph(Cora())
 GNNGraph:
diff --git a/GNNGraphs/src/temporalsnapshotsgnngraph.jl b/GNNGraphs/src/temporalsnapshotsgnngraph.jl
index a08d069a2..362689650 100644
--- a/GNNGraphs/src/temporalsnapshotsgnngraph.jl
+++ b/GNNGraphs/src/temporalsnapshotsgnngraph.jl
@@ -14,7 +14,7 @@ The features can be passed at construction time or added later.
 # Examples
 
 ```julia
-julia> using GraphNeuralNetworks
+julia> using GNNGraphs
 
 julia> snapshots = [rand_graph(10,20) for i in 1:5];
 
@@ -78,7 +78,7 @@ Return a `TemporalSnapshotsGNNGraph` created starting from `tg` by adding the sn
 # Examples
 
 ```jldoctest
-julia> using GraphNeuralNetworks
+julia> using GNNGraphs
 
 julia> snapshots = [rand_graph(10, 20) for i in 1:5];
 
@@ -138,7 +138,7 @@ Return a [`TemporalSnapshotsGNNGraph`](@ref) created starting from `tg` by remov
 # Examples
 
 ```jldoctest
-julia> using GraphNeuralNetworks
+julia> using GNNGraphs
 
 julia> snapshots = [rand_graph(10,20), rand_graph(10,14), rand_graph(10,22)];
 
diff --git a/GNNGraphs/src/transform.jl b/GNNGraphs/src/transform.jl
index 6e0b564cd..b2a88daab 100644
--- a/GNNGraphs/src/transform.jl
+++ b/GNNGraphs/src/transform.jl
@@ -170,7 +170,7 @@ A new GNNGraph with the specified edges removed.
 
 # Example
 ```julia
-julia> using GraphNeuralNetworks
+julia> using GNNGraphs
 
 # Construct a GNNGraph
 julia> g = GNNGraph([1, 1, 2, 2, 3], [2, 3, 1, 3, 1])
@@ -275,7 +275,7 @@ A new GNNGraph with the specified nodes and all edges associated with these node
 
 # Example
 ```julia
-using GraphNeuralNetworks
+using GNNGraphs
 
 g = GNNGraph([1, 1, 2, 2, 3], [2, 3, 1, 3, 1])
 
@@ -621,17 +621,17 @@ julia> e = [10.0, 20.0, 30.0, 40.0, 50.0];
 
 julia> g = GNNGraph(s, t, w, edata = e)
 GNNGraph:
-    num_nodes = 4
-    num_edges = 5
-    edata:
-        e => (5,)
+  num_nodes: 4
+  num_edges: 5
+  edata:
+        e = 5-element Vector{Float64}
 
 julia> g2 = to_bidirected(g)
 GNNGraph:
-    num_nodes = 4
-    num_edges = 7
-    edata:
-        e => (7,)
+  num_nodes: 4
+  num_edges: 7
+  edata:
+        e = 7-element Vector{Float64}
 
 julia> edge_index(g2)
 ([1, 2, 2, 3, 3, 4, 4], [2, 1, 3, 2, 4, 3, 4])