> ## Documentation Index
> Fetch the complete documentation index at: https://training-docs.cerebras.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Model

> Learn how to pass a model to the Trainer class.

The `model` is the main [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.4)") that all training and validation is run on. It's required by all [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) instances.

## Prerequisites

Read the [Trainer Overview](../../../model-zoo/trainer-overview) and [Trainer Configuration Overview](../../../model-zoo/trainer-configuration-overview) for a basic overview of how to run Model Zoo models.

## Configure the Model

Use the `model` argument to set the model you'd like to train or validate.

When using YAML, pass all `model` subkeys as arguments to the model class. Your run script's `model_fn` determines the model class.

In Python, you can specify the model in two ways:

* As a callable that takes no arguments and returns a `Module`
* As a `Module` that the system uses directly

<CodeGroup>
  ```yaml YAML theme={null}
  trainer:
    init:
      ...
      model:
        vocab_size: 1024
        max_position_embeddings: 1024
        ...
      ...
    ...

  ```

  ```python Python theme={null}
  from cerebras.modelzoo import Trainer
  from cerebras.modelzoo.models.nlp.gpt2.model import Gpt2Model

  trainer = Trainer(
      ...,
      model=lambda: Gpt2Model(
          vocab_size=1024,
          max_position_embeddings=1024,
          ...,
      ),
      ...,
  )
  ...
  ```
</CodeGroup>

<Note>
  When passing a `Module` directly, initialize the model inside the Cerebras device context for optimal performance:

  ```python theme={null}
  import cerebras.pytorch as cstorch
  from cerebras.modelzoo import Trainer
  from cerebras.modelzoo.models.nlp.gpt2.model import Gpt2Model

  # Initialize the Cerebras backend for efficient processing.
  backend = cstorch.backend("CSX")

  # Use the backend's device context manager for initializing the model.
  with backend.device:
      model = Gpt2Model(
          vocab_size=1024,
          max_position_embeddings=1024,
          ...,
      )

  # Compile the model using the Cerebras backend for optimized execution.
  trainer = Trainer(
      ...,
      backend=backend,
      model=model,
      ...,
  )
  ...
  ```

  This approach automatically moves model parameters to the Cerebras device, optimizing memory usage and improving initialization speed. For more information, see [Efficient weight initialization](../../../cs-torch/writing-a-custom-training-loop/efficient-weight-initialization).
</Note>