class FwdOnlyIdentity(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # pass the input feature map unmodified, as Identity
        return input
    
    @staticmethod
    def backward(ctx, grad_output):
        zero_grad_input = torch.zeros_like(grad_output) # set gradient w.r.t. input to 0's
        return zero_grad_input

Above is that I implemented for what I originally asked about.
Then, instead of directly doing a @ b, I added self.fwd_identity = FwdOnlyIdentity.apply. Then in the forward function, the computation becomes self.fwd_identity(a) @ b. Does this make sense?