> ## 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.

# Step Closures

By design, in the execution schema used by the Cerebras Wafer-Scale cluster, the client and the server run asynchronously to each other. This was to prevent the server from becoming bottlenecked by any client processes such as disk I/O or networking.

However, this means that a computed tensor may not be available to fetch from the server when the client requests it. For example, the call to compile happens in the first iteration of the training loop. Until compile is complete and execution starts on the cluster, no tensor is available to fetch.

To handle this, we introduce the concept of a step closure via the [`step_closure`](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/cstorch.html#cerebras.pytorch.step_closure "cerebras.pytorch.step_closure") decorator, e.g.

```python theme={null}
@cstorch.step_closure
def closure(loss):
    print(f"Loss: {loss}")
```

Tensors that are passed into a “step closure” are fetched from the server and their value is materialized before the closure is actually called. If the tensor is not yet available; it waits until the server “catches up” with the current step, and the tensor value is available to be fetched before actually calling the closure.

### Example Usage

You can call step closures either inside the traced step function or outside of it:

```python theme={null}
@cstorch.step_closure
def check_nan(loss):
    assert not torch.isnan(loss).any()

@cstorch.step_closure
def print_loss(loss):
    print(f"Loss: {loss}")

@cstorch.trace
def training_step(inputs, targets):
     outputs = compiled_model(inputs)
     loss = loss_fn(outputs, targets)

     check_nan(loss)

     loss.backward()
     optimizer.step()
     optimizer.zero_grad()

     return loss

for inputs, targets in executor:
    loss = training_step(inputs, targets)
    print_loss(loss)
```
