I’d like to parametrize my torch.jit.script'ed function with a function argument, i.e. whether a Forward-Backward algorithm should use lambda x: torch.max(x, dim = dim).values or torch.logsumexp(x, dim = dim). When I pass it as lambda, TorchScript complains that I’m calling a tensor-typed value which happens because it types the argument as Tensor (despite the fact that it’s initialized to a lambda by default).
Is there any way to tell TorchScript to type a value as a Callable?
Is it possible to template a TorchScript function by an external callable values via some other mechanism?
Should I create an issue about this on GitHub?