pytorch-minimize
pytorch-minimize copied to clipboard
make `minimize` jit-able?
Thanks for this great package, it seems thatminimize
is not jit-able? for example
from torchmin import minimize
import torch
from functools import partial
def bar(y, x):
return (x + y).square().sum()
# raise error
@torch.jit.script
def foo(x: torch.Tensor) -> torch.Tensor:
result = minimize(partial(bar, x=x), torch.tensor([0., 1.]), method='bfgs')
return result.x
Will you consider making minimize
jit-able?
Hi @alibool - thanks for your comment and suggestion.
This would require a significant refactor. I'd have to step back through each minimization routine, add type hints to the function signatures and make sure that all torch ops are TorchScript-supported. It's on my agenda but it won't be a high priority (unless I have some help).