Hello, I have the following code:
def normalize(data: torch.Tensor, mean: torch.Tensor,
              std: torch.Tensor) -> torch.Tensor:

    """Normalise the image with channel-wise mean and standard deviation.

    Args:
        data (torch.Tensor): The image tensor to be normalised.
        mean (torch.Tensor): Mean for each channel.
        std (torch.Tensor): Standard deviations for each channel.

        Returns:
        Tensor: The normalised image tensor.
    """

    if not torch.is_tensor(data):
        raise TypeError('data should be a tensor. Got {}'.format(type(data)))

    if not torch.is_tensor(mean):
        raise TypeError('mean should be a tensor. Got {}'.format(type(mean)))

    if not torch.is_tensor(std):
        raise TypeError('std should be a tensor. Got {}'.format(type(std)))

    if len(mean) != data.shape[-3] and mean.shape[:2] != data.shape[:2]:
        raise ValueError('mean lenght and number of channels do not match')

    if len(std) != data.shape[-3] and std.shape[:2] != data.shape[:2]:
        raise ValueError('std lenght and number of channels do not match')

    if std.shape != mean.shape:
        raise ValueError('std and mean must have the same shape')

    mean = mean[..., :, None, None].to(data.device)
    std = std[..., :, None, None].to(data.device)

    out = data.sub(mean).div(std)

    return out

I would like for this function to be able to be executed with and without JIT.  The problem I am currently facing is that, since this function has control flow statements (IFs) when I run:
        f = image.Normalize(mean, std)
        jit_trace = torch.jit.trace(f, data)
        jit_trace(data2)

It will run correctly, but if the input changes in such a way that the control flow takes a different branch then it will not raise the desired exception or give the desired output.
How can I make this work?  Thanks in advance!