Well, the rename to ctx is a good idea, but really, you would need to find a source for your shape.
For example TorchVision’s roi align-function takes some more parameters (vision/roi_align_kernel.cpp at 0013d9314cf1bd83eaf38c3ac6e0e9342fa99683 · pytorch/vision · GitHub 2), maybe the forward should, too, and then assign them to ctx members.