fix:orb squeeze provides incorrect shape for energy tensor

[thomasloux]

Summary

OrbModel wrapper returns a dim-0 tensor for energy for a input containing a single system, because of a squeeze() in orb.py line 419. This is not the behaviour of other models like SevenNet. This causes ts.integrate to fail when trying to split the system.

from ase.build import bulk
from torch_sim.models.orb import OrbModel
from torch_sim.state import initialize_state
from orb_models.forcefield import pretrained
import torch
import os
import torch_sim as ts
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
orbff = pretrained.orb_v3_conservative_inf_omat(
    device=device,
    precision="float32-highest",
    compile=False
)
orb_model = OrbModel(model=orbff, dtype=torch.get_default_dtype())

from torch_sim.integrators.nvt import nvt_langevin

n_steps = 30
timestep = 0.001
temperature = 300.0  # K
filename = "nvt_trajectory.h5"
save_dir = "notebook/results2"

filename = os.path.join(save_dir, filename)
reporter = ts.TrajectoryReporter(
    filenames=[filename],
    state_frequency=10,
)

system = bulk('Ti')
system = initialize_state(system, device=device, dtype=torch.get_default_dtype())
print(orb_model(system)["energy"]) # tensor(-15.5739, device='cuda:0', grad_fn=<SqueezeBackward0>)

# Run NVT simulation
_ = ts.integrate(
    system,
    orb_model,
    # sevennet_model,
    integrator=nvt_langevin,
    n_steps=n_steps,
    timestep=timestep,
    temperature=temperature,
    trajectory_reporter=reporter,
)
## RuntimeError: split expects at least a 1-dimensional tensor

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

We highly recommended installing the pre-commit hooks running in CI locally to speedup the development process. Simply run pip install pre-commit && pre-commit install to install the hooks which will check your code before each commit.

[CompRhys]

Thanks for the report! probably means we need to add another check to the model interface test to catch this potential issue

[thomasloux]

Do you any idea why the tests are failing for graphpes? It's normally unrelated to my commit.

[CompRhys]

Some of the models are flaky at the fp tolerances of the tests. I will take a look once in the office but I think it's likely you're right and its unrelated