Hi,
That does sounds like the right strategy.
In pseudo code, it would look like:
@staticmethod
def forward(ctx, inp, weight):
  ctx.save_for_backward(inp, weight)
  return cudnn_convolution(inp, weight)

@staticmethod
def backward(ctx, grad_out):
    inp, weight = ctx.saved_tensors
    return cudnn_convolution_backward_input(grad_out, weight), cudnn_convolution_backward_weight(grad_out, inp)

Doesn’t that work for you?