AttributeError: module 'habana_frameworks.torch.hpu' has no attribute 'wrap_in_hpu_graph

Training PyTorch
1

The Sentence Transformers module tries to call habana_frameworks.torch.hpu.wrap_in_hpu_graph, which seemingly does not exist.

See the error messages here.

  • I use the docker image vault.habana.ai/gaudi-docker/1.18.0/ubuntu22.04/habanalabs/pytorch-installer-2.4.0:latest and then pip install torch_geometric
  • Specifically, torch-geometric==2.6.1

See the code here

2

Can you try replacing it with:
from habana_frameworks.torch.hpu import wrap_in_hpu_graph