diff --git a/Project.toml b/Project.toml
index bbe2f84..678bca1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -24,19 +24,19 @@ WordTokenizers = "796a5d58-b03d-544a-977e-18100b691f6e"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-BSON = "0.3.3"
-CUDA = "3"
+BSON = "0.3"
+CUDA = "3, 4, 5"
 CorpusLoaders = "0.3"
 DataDeps = "0.7"
-DataStructures = "0.18.9"
-Flux = "0.12.8"
-JSON = "0.21.1"
-Languages = "0.4.3"
-NNlib = "0.7"
-StatsBase = "0.33.6"
-TextAnalysis = "0.7.3"
-WordTokenizers = "0.5.6"
-Zygote = "0.6.10"
+DataStructures = "0.18"
+Flux = "0.13, 0.14"
+JSON = "0.21"
+Languages = "0.4"
+NNlib = "0.7, 0.8, 0.9"
+StatsBase = "0.33, 0.34"
+TextAnalysis = "0.7, 0.8"
+WordTokenizers = "0.5"
+Zygote = "0.6"
 julia = "1.6"
 
 [extras]
diff --git a/src/ULMFiT/custom_layers.jl b/src/ULMFiT/custom_layers.jl
index d83c43c..2f456d9 100644
--- a/src/ULMFiT/custom_layers.jl
+++ b/src/ULMFiT/custom_layers.jl
@@ -8,7 +8,7 @@ This file contains the custom layers defined for this model:
     PooledDense
 """
 
-import Flux: gate, testmode!, _dropout_kernel
+import Flux: gate, testmode!, dropout
 
 reset_masks!(entity) = nothing
 reset_probability!(entity) = nothing
@@ -25,11 +25,11 @@ It can be used to generate the mask by giving the shape of the desired mask and
 function drop_mask(x, p)
     y = similar(x, size(x))
     Flux.rand!(y)
-    y .= Flux._dropout_kernel.(y, p, 1 - p)
+    y .= Flux.dropout(y, p)
     return y
 end
 
-drop_mask(shape::Tuple, p; type = Float32) = (mask = rand(type, shape...);mask .= _dropout_kernel.(mask, p, 1 - p))
+drop_mask(shape::Tuple, p; type = Float32) = (mask = rand(type, shape...);mask .= Flux.dropout(mask, p))
 
 #################### Weight-Dropped LSTM Cell ######################
 """
@@ -405,5 +405,5 @@ function get_trainable_params(layers)
             push!(p, l)
         end
     end
-    return params(p...)
+    return Flux.params(p...)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 2738bfa..5df3203 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,11 +2,32 @@ using Test
 using TextAnalysis
 using TextModels
 
-println("Running tests:")
-
-include("crf.jl")
-include("ner.jl")
-include("pos.jl")
-include("sentiment.jl")
-include("averagePerceptronTagger.jl")
-include("ulmfit.jl")
+ENV["DATADEPS_ALWAYS_ACCEPT"] = true
+
+tests = [
+    "crf.jl"
+    "ner.jl"
+    "pos.jl"
+    "sentiment.jl"
+    "averagePerceptronTagger.jl"
+    "ulmfit.jl"
+]
+
+function run_tests()
+    for test in tests
+        @info "Test: $test"
+        Test.@testset verbose = true "\U1F4C2 $test" begin
+            include(test)
+        end
+    end
+end
+
+@static if VERSION >= v"1.7"
+    Test.@testset verbose = true showtiming = true "All tests" begin
+        run_tests()
+    end
+else
+    Test.@testset verbose = true begin
+        run_tests()
+    end
+end
diff --git a/test/ulmfit.jl b/test/ulmfit.jl
index 3deca62..f51b685 100644
--- a/test/ulmfit.jl
+++ b/test/ulmfit.jl
@@ -1,5 +1,9 @@
+using Test
 using DataDeps
 using BSON
+using Flux
+using Flux: params
+using TextModels: ULMFiT
 
 @testset "Custom layers" begin
     @testset "WeightDroppedLSTM" begin