Merge 2becebb into 36c35ff

FluxML · Nov 12, 2020 · 1db8ca4 · 1db8ca4
2 parents 36c35ff + 2becebb
commit 1db8ca4
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git 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.1.4"
+version = "0.1.5"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"

diff --git a/README.md b/README.md
@@ -203,7 +203,9 @@ All models share the following hyper-parameters:
 
 7. `alpha`: The L2/L1 mix of regularization. Default = 0. Range = [0, 1]
 
-8. `optimiser_changes_trigger_retraining`: True if fitting an
+8. `acceleration`: Use `CUDALibs()` for training on GPU; default is `CPU1()`. 
+
+9. `optimiser_changes_trigger_retraining`: True if fitting an
    associated machine should trigger retraining from scratch whenever
    the optimiser changes. Default = `false`
 

diff --git a/src/MLJFlux.jl b/src/MLJFlux.jl
@@ -1,5 +1,7 @@
 module MLJFlux
 
+export CUDALibs, CPU1
+
 import Flux
 using MLJModelInterface
 using MLJModelInterface.ScientificTypes