I don’t understand why you want max to be saved on the Module?
And why you cannot just pass it along?
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):
      ...
      EDIT: need to return None for non-differentiable inputs
      return grad_input, None
  
  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)
    output = self.uni(x, max)
    return output