Closed peastman closed 2 years ago
Hi, Thanks for letting us know. As you guys mentioned in another thread, this should be a bug for pytorch version < 1.8.
@yueyericardo @peastman yes, thank you, we are aware of the issues with JIT, this is actually still an issue in > 1.8, specifically it is also an issue in nightly right now if you load the model from C++, I have a workaround for it that I'm using in a private repo but I may merge it here. basically the issue is related to TorchScripts parsing of
pair_sizes = counts * (counts - 1) // 2
for some reason TorchScript creates buggy code that sometimes executes (counts ** 2) // 2 by making 1 into a tensor the issue can be bypassed with no performance penalty
one = torch.tensor(1, dtype=torch.long, device=counts.device)
pair_sizes = counts * (counts - one) // 2
it is ugly but it works
TorchANI models do not work correctly when they are converted to TorchScript. I'm pretty sure this is a PyTorch bug, but you should be aware of it, and perhaps you can find a workaround. The following script demonstrates the problem using PyTorch 1.7.1.
It fails with this exception.
Important points:
I tried to debug it a bit further, and I believe the error is happening in this line:
https://github.com/aiqm/torchani/blob/2007d181c1d650e28ee4e2c6a5f8e74fae54d6bf/torchani/aev.py#L249
I print out both
counts
andpair_sizes
, and I can see that the second time it is invoked,pair_sizes
is computed incorrectly:For example, the first element should be 10*9/2 = 45, not 49.