pad: add optional padding mode

This commit is contained in:
Cédric Deltheil 2023-10-04 17:36:04 +02:00 committed by Cédric Deltheil
parent 9b1e25e682
commit f4298f87d2

View file

@ -23,8 +23,8 @@ def manual_seed(seed: int) -> None:
_manual_seed(seed) _manual_seed(seed)
def pad(x: Tensor, pad: Iterable[int], value: float = 0.0) -> Tensor: def pad(x: Tensor, pad: Iterable[int], value: float = 0.0, mode: str = "constant") -> Tensor:
return _pad(input=x, pad=pad, value=value) # type: ignore return _pad(input=x, pad=pad, value=value, mode=mode) # type: ignore
def interpolate(x: Tensor, factor: float | torch.Size, mode: str = "nearest") -> Tensor: def interpolate(x: Tensor, factor: float | torch.Size, mode: str = "nearest") -> Tensor: