Hi!
I’m trying to load a PyTorch model in C++, using JIT. The model is defined as follows:
class JitModel(torch.jit.ScriptModule):
    def __init__(self):
        super(JitModel, self).__init__()

        self.n_layers = 5
        self.n_features = 14
        self.fc1 = torch.nn.Linear(14, 14)
        self.fc2 = torch.nn.Linear(14, 14)
        self.fc3 = torch.nn.Linear(14, 14)
        self.fc4 = torch.nn.Linear(14, 14)
        self.fc5 = torch.nn.Linear(14, 14)
        self.out = torch.nn.Linear(14, 1)
        self.normalise = torch.nn.BatchNorm1d(14)

    @torch.jit.script_method
    def forward(self, x):
        _x = x
        # _x = self.normalise(_x)
        _x = F.relu(self.fc1(_x))
        _x = F.relu(self.fc2(_x))
        _x = F.relu(self.fc3(_x))
        _x = F.relu(self.fc4(_x))
        _x = F.relu(self.fc5(_x))
        _x = torch.sigmoid(self.out(_x))
        return _x

where
F = torch.nn.functional

This model is used for training in a Python3 script and then saved using the JIT save() function.
When loaded with
torch.jit.load("model.pt") (in Python)

the model is loaded correctly.
When loaded with
torch::jit::load("model.pt") (in C++)

this happens:
terminate called after throwing an instance of 'torch::jit::script::ErrorReport'
  what():
Return value was annotated as having type Tuple[] but is actually of type Optional[Tuple[]]:
op_version_set = 0
def _check_input_dim(self,
    input: Tensor) -> Tuple[]:
  _0 = torch.ne(torch.dim(input), 2)
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~...  <--- HERE
  if _0:
    _1 = torch.ne(torch.dim(input), 3)
  else:
    _1 = _0
  if _1:
    ops.prim.RaiseException("Exception")
  else:
    pass
def forward(self,
Aborted (core dumped)

If I comment the line:
        self.normalise = torch.nn.BatchNorm1d(14)

in the definition, the C++ script loads the model correctly.
Really no idea of what’s wrong with this implementation.