Switch to CUDA implementation if batch size >= 65536 for affine_grid (#16403)

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

diff --git a/torch/nn/_functions/vision.py b/torch/nn/_functions/vision.py
index 159025e..018e93b 100644 (file)
--- a/torch/nn/_functions/vision.py
+++ b/torch/nn/_functions/vision.py
@@ -6,7 +6,7 @@ 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:
+    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: