Where is torch.cudnn_affine_grid_generator?

(Xinke) #1

I want to modify the affine_grid_generator function in torch/nn/_functions/vision.py from 2D to 3D. But I can’t find where is torch.cudnn_affine_grid_generator function. My pytorch version is 0.4.0 and python version is 3.6. The source code is shown below. Thanks in advance!

def affine_grid_generator(theta, size):
if theta.data.is_cuda:
if not cudnn.enabled:
raise RuntimeError("AffineGridGenerator needs CuDNN for "
“processing CUDA inputs, but CuDNN is not enabled”)
if not cudnn.is_acceptable(theta.data):
raise RuntimeError(“AffineGridGenerator generator theta not acceptable for CuDNN”)
N, C, H, W = size
return torch.cudnn_affine_grid_generator(theta, N, C, H, W)
return AffineGridGenerator.apply(theta, size)

def affine_grid_generator_3d(theta, size):
return AffineGridGenerator_3d.apply(theta, size)

(Thomas V) #2

It’s in the C++ modules (ATen to be precise).

Personally, I tend to use the “rgrep” command on the PyTorch source tree to find functions.

Best regards


(Mohamed Hassan) #3

Hi @Keke
Did you solve this issue. I want to do exactly the same.

(Teaghan) #4

I am also interested in a solution to this. @Keke did you happen to build a 3D affine grid generator?