Python PyTorch wrap_torch_function用法及代碼示例

本文簡要介紹python語言中 torch.overrides.wrap_torch_function 的用法。

用法:

torch.overrides.wrap_torch_function(dispatcher)

參數：

dispatcher(可調用) -返回傳遞給函數的 Tensor-likes 的可迭代對象的可調用對象。

使用與 __torch_function__ 相關的函數包裝給定函數。

注意

此裝飾器可能會降低代碼的性能。通常，將您的代碼表達為一係列本身支持 __torch_function__ 的函數就足夠了。如果您發現自己處於罕見的情況並非如此，例如如果您正在包裝一個低級庫並且您還需要它為Tensor-likes 工作，那麽這個函數是可用的。

例子

>>> def dispatcher(a): # Must have the same signature as func
...     return (a,)
>>> @torch.overrides.wrap_torch_function(dispatcher)
>>> def func(a): # This will make func dispatchable by __torch_function__
...     return a + 0