From 4bf276d395479146a1ba7bffe9f06a99f54d3fbe Mon Sep 17 00:00:00 2001 From: giacomo m Date: Sun, 21 Apr 2024 12:02:29 +0000 Subject: [PATCH] Don't enter inference mode in solver. It's redundant with falkon.fit --- falkon/optim/conjgrad.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/falkon/optim/conjgrad.py b/falkon/optim/conjgrad.py index a6b8157..2ea425f 100644 --- a/falkon/optim/conjgrad.py +++ b/falkon/optim/conjgrad.py @@ -282,7 +282,7 @@ def solve(self, X, M, Y, _lambda, initial_solution, max_iter, callback=None): stream = torch.cuda.current_stream(device) # Note that if we don't have CUDA this still works with stream=None. - with ExitStack() as stack, TicToc("ConjGrad preparation", False), torch.inference_mode(): + with ExitStack() as stack, TicToc("ConjGrad preparation", False): if cuda_inputs: stack.enter_context(torch.cuda.device(device)) stack.enter_context(torch.cuda.stream(stream))