I come across the same problem. Below is my code:
class Multipler(torch.autograd.Function):
  """
  We can implement our own custom autograd Functions by subclassing
  torch.autograd.Function and implementing the forward and backward passes
  which operate on Tensors.
  """
  def forward(self, input):
    """
    In the forward pass we receive a Tensor containing the input and return a
    Tensor containing the output. You can cache arbitrary Tensors for use in the
    backward pass using the save_for_backward method.
    """
    self.save_for_backward(input)
    output = input * 10000
    return output

  def backward(self, grad_output):
    """
    In the backward pass we receive a Tensor containing the gradient of the loss
    with respect to the output, and we need to compute the gradient of the loss
    with respect to the input.
    """
    input, = self.saved_tensors
    grad_input = grad_output.clone()
    return grad_input

class normalize(torch.nn.Module):
    def __init__(self):
        super(normalize, self).__init__()
        self.Softmax_layer1 = nn.Softmax2d()
        self.Softmax_layer2 = nn.Softmax2d()
        
    def forward(self, inputs):
        output = self.Softmax_layer1(Multipler.apply(inputs))
        output = self.Softmax_layer2(Multipler.apply(output))
        
        return output

But when I delete the *10000 in the forward function in the class Multipler, there is no such   error. But what I expect is the mismatch between the forward and backward. That is my design. For my case how to solve this problem?