mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
82175f31b4
Summary: Port AffineGrid to C++, because script does not support compiling Function classes. Pull Request resolved: https://github.com/pytorch/pytorch/pull/14392 Differential Revision: D13219698 Pulled By: eellison fbshipit-source-id: 3ddad8a84c72010b5a6c6f7f9712be614202faa6
14 lines
430 B
Python
14 lines
430 B
Python
import torch
|
|
import torch.backends.cudnn as cudnn
|
|
|
|
|
|
@torch._jit_internal.weak_script
|
|
def affine_grid_generator(theta, size):
|
|
# type: (Tensor, List[int]) -> Tensor
|
|
if theta.is_cuda and cudnn.enabled and cudnn.is_acceptable(theta) and len(size) == 4:
|
|
N, C, H, W = size
|
|
ret = torch.cudnn_affine_grid_generator(theta, N, C, H, W)
|
|
else:
|
|
ret = torch.affine_grid_generator(theta, size)
|
|
return ret
|