Model loads correctly on Windows 11 but not Linux

jit
1

Hi all,

I’m attempting to load a model with torch.JIT.load(‘model.pt’) and the model correctly loads on Windows 11. However, if I load the model on Linux, it fails with the following error in the forward() call:

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The size of tensor a (335) must match the size of tensor b (3) at non-singleton dimension 0

It doesn’t matter what OS the model was saved on, it always fails. I have no idea where to even start looking to fix this issue, any help would be appreciated!

2

Could you check which line of code fails exactly with this error? I don’t know why the shape mismatch would be raised on Linux but not on Windows though.

3

energy = self.model((species, coordinates), cell=cell, pbc=pbc)

4

This is where the error is thrown in pytorch

"/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
5

Can you run with the environment variable TORCH_SHOW_CPP_STACKTRACES=1 set?

6 
  File "/home/ray/Dropbox/repos/LightningANI/lightning_ani/anscan.py", line 47, in <module>
    dyn.run()
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/optimize/optimize.py", line 269, in run
    return Dynamics.run(self)
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/optimize/optimize.py", line 156, in run
    for converged in Dynamics.irun(self):
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/optimize/optimize.py", line 135, in irun
    self.step()
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/optimize/bfgslinesearch.py", line 122, in step
    ls._line_search(self.func, self.fprime, r, self.p, g, e, self.e0,
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/utils/linesearch.py", line 57, in _line_search
    fval = func(xk + stp * pk, *args)
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/optimize/bfgslinesearch.py", line 171, in func
    return (self.atoms.get_potential_energy(
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/atoms.py", line 728, in get_potential_energy
    energy = self._calc.get_potential_energy(
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/calculators/calculator.py", line 709, in get_potential_energy
    energy = self.get_property('energy', atoms)
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/ase/calculators/calculator.py", line 737, in get_property
    self.calculate(atoms, [name], system_changes)
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/torchani/ase.py", line 80, in calculate
    energy = self.model((species, coordinates), cell=cell, pbc=pbc)
  File "/home/ray/anaconda3/envs/torchani/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The size of tensor a (31080) must match the size of tensor b (3) at non-singleton dimension 0