@ptrblck Thanks for your reply. I actually figured out the right way to do this. I basically wrote my own Autograd Function class, similar to here: https://pytorch.org/docs/stable/notes/extending.html 17
Then inside its backward method, I used a different set of weights to compute grad_input. So the backward now looks like this:
    def backward(ctx, grad_output):
        input, weight, b_weights, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(b_weights)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[3]:
            grad_bias = grad_output.sum(0).squeeze(0)

        return grad_input, grad_weight, torch.zeros_like(b_weights), grad_bias

b_weights are the backwards weights that are passed to the forward function and saved in ctx