Unable to Use Apple GPU (MPS) with OCPCalculator

[notosaka]

Hello FairChem team,

I am trying to run a molecular dynamics simulation using the FairChem library on my Apple Silicon device. Specifically, I want to leverage the Apple GPU (MPS) for the calculations. Here is the code I am using:

import torch
from ase.build import fcc100, add_adsorbate, molecule
from ase.optimize import LBFGS
from fairchem.core import OCPCalculator

# Check if MPS (Apple GPU) is available
print(torch.backends.mps.is_available())  # Verifies Metal Performance Shaders (MPS) availability
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

# Set up your system as an ASE atoms object
slab = fcc100('Cu', (5, 5, 5), vacuum=8)
adsorbate = molecule("CO")
add_adsorbate(slab, adsorbate, 2.0, 'bridge')

# OCPCalculator settings (using GPU)
calc = OCPCalculator(
    model_name="EquiformerV2-31M-S2EF-OC20-All+MD",
    local_cache="pretrained_models",
    cpu=False  # Set to False to use Apple GPU (MPS)
)
slab.calc = calc

# Set up LBFGS dynamics object
dyn = LBFGS(slab)
dyn.run(0.05, 100)

However, despite setting cpu=False in the OCPCalculator, the calculations seem to be running on the CPU rather than the GPU. I checked the GPU usage via the Activity Monitor, and it shows 0% usage, which indicates that the GPU is not being utilized.

Could you please provide guidance on how to properly configure OCPCalculator to utilize the Apple GPU (MPS)? Am I missing any crucial steps or configurations?

Thank you for your assistance.

Best regards,
Osaka

[github-actions]

This issue has been marked as stale because it has been open for 30 days with no activity.