mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
f3cc9b2218
Summary: Cleanup to make references to `weak_script` consistent across codebase Pull Request resolved: https://github.com/pytorch/pytorch/pull/15364 Differential Revision: D13509676 Pulled By: driazati fbshipit-source-id: 93dbbbe57e9b9b6587895f3cc6fac678babd21de
15 lines
451 B
Python
15 lines
451 B
Python
import torch
|
|
import torch.backends.cudnn as cudnn
|
|
from ..._jit_internal import weak_script
|
|
|
|
|
|
@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
|