diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl index 1b856dc..e952b7b 100644 --- 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 index 0f0e9ca..ee8b405 100644 --- 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,