Merge branch 'master' of github.com:FalkonML/falkon

FalkonML · May 1, 2024 · b4c7814 · b4c7814
2 parents 51d0885 + 4bf276d
commit b4c7814
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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))