Thanks for the replies. However I cannot find the environment variable PYTORCH_JIT in my system. To search for environment variables I use the following code:
for key in os.environ.keys():
      if "py" in key.lower():
          print(key)

Which yields:
CONDA_PYTHON_EXE
PYTHONPATH
Nothing Pytorch related is shown.
On a second note. The @torch.script.jit runs the jitted code and throws the following error.
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/diego/Projects/torchgeometry/.dev_env/lib/python3.7/site-packages/torch/nn/mo
dules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/diego/Projects/torchgeometry/torchgeometry/image/normalization.py", line 28,
in forward
    return normalize(input, self.mean, self.std)
RuntimeError:
Dimension out of range (expected to be in range of [-1, 0], but got -3) (maybe_wrap_dim at
/opt/conda/conda-bld/pytorch_1556653215914/work/c10/core/WrapDimMinimal.h:20)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x45 (0x7ff3515d7dc5
 in /home/diego/Projects/torchgeometry/.dev_env/lib/python3.7/site-packages/torch/lib/libc1
0.so)
frame #1: <unknown function> + 0x7d4a40 (0x7ff351fc5a40 in /home/diego/Projects/torchgeomet
ry/.dev_env/lib/python3.7/site-packages/torch/lib/libcaffe2.so)
frame #2: at::native::slice(at::Tensor const&, long, long, long, long) + 0x4e (0x7ff351fc63
8e in /home/diego/Projects/torchgeometry/.dev_env/lib/python3.7/site-packages/torch/lib/lib
caffe2.so)
frame #3: at::TypeDefault::slice(at::Tensor const&, long, long, long, long) const + 0x1a (0
x7ff3522255fa in /home/diego/Projects/torchgeometry/.dev_env/lib/python3.7/site-packages/to
rch/lib/libcaffe2.so)
frame #4: torch::autograd::VariableType::slice(at::Tensor const&, long, long, long, long) c
onst + 0x6d3 (0x7ff34a32ea23 in /home/diego/Projects/torchgeometry/.dev_env/lib/python3.7/s
ite-packages/torch/lib/libtorch.so.1)
frame #5: <unknown function> + 0x985117 (0x7ff34a4ec117 in /home/diego/Projects/torchgeomet
ry/.dev_env/lib/python3.7/site-packages/torch/lib/libtorch.so.1)
frame #6: <unknown function> + 0xa73df8 (0x7ff34a5dadf8 in /home/diego/Projects/torchgeomet
ry/.dev_env/lib/python3.7/site-packages/torch/lib/libtorch.so.1)
frame #7: torch::jit::InterpreterState::run(std::vector<c10::IValue, std::allocator<c10::IV
alue> >&) + 0x22 (0x7ff34a5d6372 in /home/diego/Projects/torchgeometry/.dev_env/lib/python3
.7/site-packages/torch/lib/libtorch.so.1)
frame #8: <unknown function> + 0xa5b2d9 (0x7ff34a5c22d9 in /home/diego/Projects/torchgeomet
ry/.dev_env/lib/python3.7/site-packages/torch/lib/libtorch.so.1)
frame #9: <unknown function> + 0x458bf3 (0x7ff3776d0bf3 in /home/diego/Projects/torchgeomet
ry/.dev_env/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #10: <unknown function> + 0x12d07a (0x7ff3773a507a in /home/diego/Projects/torchgeome
try/.dev_env/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #37: __libc_start_main + 0xe7 (0x7ff384563b97 in /lib/x86_64-linux-gnu/libc.so.6)
:
operation failed in interpreter:
        raise ValueError('mean lenght and number of channels do not match')

    if std.shape[0] != 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)
           ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    std = std[..., :, None, None].to(data.device)

    out = (data - mean) / std

    return out

What am I doing wrong?
I am using Pytorch 1.1 on an Anaconda envrionment