fix LRFinder

OHDSI · Apr 19, 2023 · f26382e · f26382e
1 parent 935f51d
commit f26382e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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)