diff --git a/Project.toml b/Project.toml
index cac18ad2..589216d8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "MLJFlux"
 uuid = "094fc8d1-fd35-5302-93ea-dabda2abf845"
 authors = ["Anthony D. Blaom <anthony.blaom@gmail.com>", "Ayush Shridhar <ayush.shridhar1999@gmail.com>"]
-version = "0.3.1"
+version = "0.4.0"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
@@ -19,22 +19,25 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 CategoricalArrays = "0.10"
 ColorTypes = "0.10.3, 0.11"
 ComputationalResources = "0.3.2"
-Flux = "0.13, 0.14"
+Flux = "0.14"
 MLJModelInterface = "1.1.1"
-Metalhead = "0.8"
+Metalhead = "0.9"
 ProgressMeter = "1.7.1"
+StatisticalMeasures = "0.1"
 Tables = "1.0"
-julia = "1.6"
+julia = "1.9"
 
 [extras]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+StatisticalMeasures = "a19d573c-0a75-4610-95b3-7071388c7541"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["cuDNN", "LinearAlgebra", "MLJBase", "Random", "StableRNGs", "Statistics", "StatsBase", "Test"]
+test = ["CUDA", "cuDNN", "LinearAlgebra", "MLJBase", "Random", "StableRNGs", "StatisticalMeasures", "Statistics", "StatsBase", "Test"]
diff --git a/test/builders.jl b/test/builders.jl
index 0f7e27a9..ba7df095 100644
--- a/test/builders.jl
+++ b/test/builders.jl
@@ -31,7 +31,7 @@ MLJFlux.build(builder::TESTBuilder, rng, n_in, n_out) =
     chain0 = myinit(1, d)
     pretraining_yhat = Xmat*chain0' |> vec
     @test y isa Vector && pretraining_yhat isa Vector
-    pretraining_loss_by_hand =  MLJBase.l2(pretraining_yhat, y) |> mean
+    pretraining_loss_by_hand =  StatisticalMeasures.l2(pretraining_yhat, y) |> mean
     mean(((pretraining_yhat - y).^2)[1:2])
 
     # compare:
diff --git a/test/classifier.jl b/test/classifier.jl
index 3eb73699..82de9500 100644
--- a/test/classifier.jl
+++ b/test/classifier.jl
@@ -56,7 +56,7 @@ losses = []
     end
     dist = MLJBase.UnivariateFinite(prob_given_class)
     loss_baseline =
-        MLJBase.cross_entropy(fill(dist, length(test)), y[test]) |> mean
+        StatisticalMeasures.cross_entropy(fill(dist, length(test)), y[test]) |> mean
 
     # check flux model is an improvement on predicting constant
     # distribution:
@@ -71,7 +71,7 @@ losses = []
     first_last_training_loss = MLJBase.report(mach)[1][[1, end]]
     push!(losses, first_last_training_loss[2])
     yhat = MLJBase.predict(mach, rows=test);
-    @test mean(MLJBase.cross_entropy(yhat, y[test])) < 0.95*loss_baseline
+    @test mean(StatisticalMeasures.cross_entropy(yhat, y[test])) < 0.95*loss_baseline
 
     optimisertest(MLJFlux.NeuralNetworkClassifier,
                   X,
diff --git a/test/integration.jl b/test/integration.jl
index 70b45fd1..39f4f79c 100644
--- a/test/integration.jl
+++ b/test/integration.jl
@@ -12,13 +12,13 @@ y, X = unpack(table, ==(:target), _->true, rng=rng)
     model3 = deepcopy(model)
     model3.lambda = 0.1
 
-    e = evaluate(model, X, y, resampling=Holdout(), measure=LogLoss())
+    e = evaluate(model, X, y, resampling=Holdout(), measure=StatisticalMeasures.LogLoss())
     loss1 = e.measurement[1]
 
-    e = evaluate(model2, X, y, resampling=Holdout(), measure=LogLoss())
+    e = evaluate(model2, X, y, resampling=Holdout(), measure=StatisticalMeasures.LogLoss())
     loss2 = e.measurement[1]
 
-    e = evaluate(model3, X, y, resampling=Holdout(), measure=LogLoss())
+    e = evaluate(model3, X, y, resampling=Holdout(), measure=StatisticalMeasures.LogLoss())
     loss3 = e.measurement[1]
 
     @test loss1 ≈ loss2
diff --git a/test/mlj_model_interface.jl b/test/mlj_model_interface.jl
index a2e6c8fe..52452ea5 100644
--- a/test/mlj_model_interface.jl
+++ b/test/mlj_model_interface.jl
@@ -30,7 +30,7 @@ end
     @test model.batch_size == 1
 
     if MLJFlux.gpu_isdead()
-        model  = @test_logs (:warn, r"`acceleration") begin
+        model  = @test_logs (:warn, r"`acceleration") match_mode = :any begin
             ModelType(acceleration=CUDALibs())
         end
         @test model.acceleration == CUDALibs()
diff --git a/test/runtests.jl b/test/runtests.jl
index 691edb89..d8df35be 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -10,7 +10,8 @@ import Random.seed!
 using Statistics
 import StatsBase
 using StableRNGs
-using cuDNN
+using CUDA, cuDNN
+import StatisticalMeasures
 
 using ComputationalResources
 using ComputationalResources: CPU1, CUDALibs