From: vishwakftw <cs15btech11043@iith.ac.in>
Date: Sat, 26 Jan 2019 19:14:19 +0000 (-0800)
Subject: Switch to CUDA implementation if batch size >= 65536 for affine_grid (#16403)
X-Git-Tag: accepted/tizen/6.5/unified/20211028.231830~1663
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=8c81a72e87ac4114fbddfadb9285d89a9ead79a1;p=platform%2Fupstream%2Fpytorch.git

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
--- 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: