From f26382e34c85fd22cf2c1a253efae4c2d6b9c9a3 Mon Sep 17 00:00:00 2001
From: egillax <egillax@gmail.com>
Date: Wed, 19 Apr 2023 17:57:29 +0200
Subject: [PATCH] fix LRFinder

---
 R/LRFinder.R | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/R/LRFinder.R b/R/LRFinder.R
index 2e79fca..2ef7d8c 100644
--- a/R/LRFinder.R
+++ b/R/LRFinder.R
@@ -52,7 +52,10 @@ lrFinder <- function(dataset, modelType, modelParams, estimatorSettings,
                      divergenceThreshold=4) {
     torch::torch_manual_seed(seed=estimatorSettings$seed)
     model <- do.call(modelType, modelParams)
-    model$to(device=estimatorSettings$device)
+    if (is.function(estimatorSettings$device)) {
+      device = estimatorSettings$device()
+    } else {device = estimatorSettings$device}
+    model$to(device=device)
     
     optimizer <- estimatorSettings$optimizer(model$parameters, lr=minLR)