def _find_module_of_method(orig_method: Callable[(..., Any)]) -> str:
    name = orig_method.__name__
    module = orig_method.__module__
    if (module is not None):
        return module
    for guess in [torch, torch.nn.functional]:
        if (getattr(guess, name, None) is orig_method):
            return guess.__name__
    raise RuntimeError(f'cannot find module for {orig_method}')
def _find_module_of_method(orig_method: Callable[(..., Any)]) -> str:
    name = orig_method.__name__
    module = orig_method.__module__
    if (module is not None):
        return module
    for guess in [torch, torch.nn.functional]:
        if (getattr(guess, name, None) is orig_method):
            return guess.__name__
    raise RuntimeError(f'cannot find module for {orig_method}')