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)