Cerebras PyTorch API#

The cerebras_pytorch package contains functions and data structures that can be used to configure and execute PyTorch models on a Cerebras Wafer Scale cluster.

cerebras_pytorch.amp

Module for configuring and using automatic mixed precision

cerebras_pytorch.optim

Module containing Cerebras compliant optimizers

cerebras_pytorch.sparse

Module containing Cerebras sparsity routines

cerebras_pytorch.metrics

Module containing Cerebras compliant metrics