Porting PyTorch Models to Gaudi

This section details the steps required to prepare PyTorch models to run on Gaudi. To create and execute a simple PyTorch model with Gaudi, see Getting Started with PyTorch and Gaudi.

  1. Import Habana Torch Library:

import habana_frameworks.torch.core as htcore

  1. Target the Gaudi HPU device:

device = torch.device("hpu")

  1. Add mark_step(). In Lazy mode, mark_step() must be added in all training scripts right after loss.backward() and optimizer.step().

htcore.mark_step()

Note

  • Placing mark_step() at any arbitrary point in the code is currently not supported.