Update to resolve FutureWarning from torch.load

lanl-ansi · Oct 20, 2024 · a9535c5 · a9535c5
1 parent 40987b5
commit a9535c5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl
@@ -102,7 +102,7 @@ function MathOptAI.build_predictor(
     end
     torch = PythonCall.pyimport("torch")
     nn = PythonCall.pyimport("torch.nn")
-    torch_model = torch.load(predictor.filename)
+    torch_model = torch.load(predictor.filename; weights_only = false)
     return _predictor(nn, torch_model, config)
 end
 
@@ -134,7 +134,7 @@ function MathOptAI.GrayBox(
     hessian::Bool = false,
 )
     torch = PythonCall.pyimport("torch")
-    torch_model = torch.load(predictor.filename)
+    torch_model = torch.load(predictor.filename; weights_only = false)
     J = torch.func.jacrev(torch_model)
     H = torch.func.hessian(torch_model)
     # TODO(odow): I'm not sure if there is a better way to get the output

diff --git a/test/test_PythonCall.jl b/test/test_PythonCall.jl
@@ -34,7 +34,7 @@ end
 
 function _evaluate_model(filename, x)
     torch = PythonCall.pyimport("torch")
-    torch_model = torch.load(filename)
+    torch_model = torch.load(filename; weights_only = false)
     input = torch.tensor(x)
     return PythonCall.pyconvert(Vector, torch_model(input).detach().numpy())
 end
@@ -299,7 +299,7 @@ function test_model_Tanh_scalar_GrayBox()
     )
     model = Model(Ipopt.Optimizer)
     set_silent(model)
-    @variable(model, x[1:2])
+    @variable(model, x[i in 1:2] == i)
     ml_model = MathOptAI.PytorchModel(filename)
     y, formulation =
         MathOptAI.add_predictor(model, ml_model, x; gray_box = true)
@@ -331,7 +331,7 @@ function test_model_Tanh_scalar_GrayBox_hessian()
     )
     model = Model(Ipopt.Optimizer)
     set_silent(model)
-    @variable(model, x[1:2])
+    @variable(model, x[i in 1:2] == i)
     ml_model = MathOptAI.PytorchModel(filename)
     y, formulation = MathOptAI.add_predictor(
         model,
@@ -369,7 +369,7 @@ function test_model_Tanh_vector_GrayBox()
     # Full-space
     model = Model(Ipopt.Optimizer)
     set_silent(model)
-    @variable(model, x[1:3])
+    @variable(model, x[i in 1:3] == i)
     ml_model = MathOptAI.PytorchModel(filename)
     y, formulation =
         MathOptAI.add_predictor(model, ml_model, x; gray_box = true)
@@ -381,7 +381,7 @@ function test_model_Tanh_vector_GrayBox()
     # Reduced-space
     model = Model(Ipopt.Optimizer)
     set_silent(model)
-    @variable(model, x[1:3])
+    @variable(model, x[i in 1:3] == i)
     ml_model = MathOptAI.PytorchModel(filename)
     y, formulation = MathOptAI.add_predictor(
         model,
@@ -419,7 +419,7 @@ function test_model_Tanh_vector_GrayBox_hessian()
     # Full-space
     model = Model(Ipopt.Optimizer)
     set_silent(model)
-    @variable(model, x[1:3])
+    @variable(model, x[i in 1:3] == i)
     ml_model = MathOptAI.PytorchModel(filename)
     y, formulation = MathOptAI.add_predictor(
         model,
@@ -436,7 +436,7 @@ function test_model_Tanh_vector_GrayBox_hessian()
     # Reduced-space
     model = Model(Ipopt.Optimizer)
     set_silent(model)
-    @variable(model, x[1:3])
+    @variable(model, x[i in 1:3] == i)
     ml_model = MathOptAI.PytorchModel(filename)
     y, formulation = MathOptAI.add_predictor(
         model,