format

lanl-ansi · Nov 5, 2024 · 1d9e8f2 · 1d9e8f2
1 parent 594911f
commit 1d9e8f2
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl
@@ -147,17 +147,15 @@ function MathOptAI.GrayBox(
     device::String = "cpu",
 )
     torch = PythonCall.pyimport("torch")
-    torch_model = torch.load(
-        predictor.filename;
-        weights_only = false,
-    ).to(device)
+    torch_model = torch.load(predictor.filename; weights_only = false)
+    torch_model = torch_model.to(device)
     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
     # dimension of a torch model object?
     output_size(::Any) = PythonCall.pyconvert(Int, torch_model[-1].out_features)
     function callback(x)
-        py_x = torch.tensor(collect(x), device = device)
+        py_x = torch.tensor(collect(x); device = device)
         py_value = torch_model(py_x).detach().cpu().numpy()
         value = PythonCall.pyconvert(Vector, py_value)
         py_jacobian = J(py_x).detach().cpu().numpy()