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)

else:

return AffineGridGenerator.apply(theta, size)

def affine_grid_generator_3d(theta, size):

return AffineGridGenerator_3d.apply(theta, size)