mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
8c81a72e87
Summary: Changelog: - Append a condition that switches to the native CUDA implementation for affine_grid Fixes #16365 Differential Revision: D13832192 Pulled By: soumith fbshipit-source-id: 3f484e6673d71e3ba7627b170cb8f1611e12b9b2
15 lines
471 B
Python
15 lines
471 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 and size[0] < 65536:
|
|
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
|