class Mapping_multi(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, num_partial_sum):
        output = input.clone()
        with torch.no_grad():
            error_bias_array_real = (torch.multinomial(error_distribution, num_partial_sum, replacement=True) - 35).float().permute(1,0)
            error_array_real = error_bias_array_real + identity_array[None, : ]
            error_array = error_array_real.float()
            value_index = (output+35).long()
            output = error_array[torch.arange(num_partial_sum)[:, None, None, None], value_index]
        return output

    @staticmethod
    def backward(ctx, grad_output):
           grad_input = grad_output.clone()
        return grad_input, None

Here is my code, if I add this function after the convolution at every layer, it will increase much time.
This function is that I want to replace the activation to error value.
error array is 2-d tensor with error value, I want replace the activation according to its value with error array.