Add more details for Callable

[Curt-Park]

e.g. Callable[[torch.Tensor], torch.Tensor]