Workaround for jit export error for - built-in method apply of FunctionMeta

jit
1 
from torch.autograd import Function

class SubMConvFunction(Function):
    @staticmethod
    def forward(ctx, features, filters, indice_pairs, indice_pair_num,
                num_activate_out, algo):
        ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
        ctx.algo = algo
        return ops.indice_conv(features,
                               filters,
                               indice_pairs,
                               indice_pair_num,
                               num_activate_out,
                               False,
                               True,
                               algo=algo)

    @staticmethod
    def backward(ctx, grad_output):
        indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
        input_bp, filters_bp = ops.indice_conv_backward(features,
                                                        filters,
                                                        grad_output,
                                                        indice_pairs,
                                                        indice_pair_num,
                                                        False,
                                                        True,
                                                        algo=ctx.algo)

        return input_bp, filters_bp, None, None, None, None

indice_subm_conv = SubMConvFunction.apply

This giving error that

Python builtin <built-in method apply of FunctionMeta object at 0x56a2498> is currently not supported in Torchscript:
if self.subm:
                out_features = Fsp.indice_subm_conv(features, self.weight,
                               ~~~~~~~~~~~~~~~~~~~~ <--- HERE
                                                    indice_pairs.to(device),
                                                    indice_pair_num,

I know that the autograd is currently not supported in the torch jit export. My model training is done, so I’m only interested in inference and exporting. Is there a way to replace this whole thing with just a simple forward function which is supported by jit.