I don’t know if pass into forward function is a good idea? I tried your solution but prompted RuntimeError: function qfnBackward returned an incorrect number of gradients (expected 2, got 1). I think the backward might consider the max_value for backprop also, so I tried to set the max_value.requires_grad = False. But the error still remained.
def uniform(k):
  class qfn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, max_value):
      ...
      out = out / max_value
      return out

    @staticmethod
    def backward(ctx, grad_output):
      ...
      return grad_input
  
  return qfn().apply

class target_fn(nn.Module):
  def __init__(self, bit):
    super(target_fn, self).__init__()
    ...
    self.uni = uniform(k=bit)

  def forward(self, x):
    max = max(x)
    max.requires_grad = False
    output = self.uni(x, max)
    return output

The reason I want max to be saved on the module is to properly construct the uniform. You see, the uniform.qfn.forward contains some logic based on the max.
Thanks again, albanD.