From 8b3f8f2942633155f51e49136398bdee178e4017 Mon Sep 17 00:00:00 2001
From: MariaHei <mabuwu83@gmail.com>
Date: Fri, 23 Aug 2024 13:34:23 +0100
Subject: [PATCH] get_and_train_model: provided model is moved to gpu if
 available

---
 src/deep_learning.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/deep_learning.jl b/src/deep_learning.jl
index 60c4866..b8ee6e3 100644
--- a/src/deep_learning.jl
+++ b/src/deep_learning.jl
@@ -136,6 +136,8 @@ function get_and_train_model(X_train::Union{SparseMatrixCSC,Matrix},
         model = Chain(
             Dense(size(X_train, 2) => hidden_dim, relu),   # activation function inside layer
             Dense(hidden_dim => size(Y_train, 2))) |> gpu        # move model to GPU, if available
+    else
+        model = model |> gpu
     end
 
     verbose && @show model