SingleDeviceStrategy

class braintools.trainer.SingleDeviceStrategy(device=None)

Single device strategy (no distribution).

This is the default strategy when only one device is available or distribution is not needed.

Parameters:

device (Any | None) – The device to use. Default is the first available device.

Examples

>>> strategy = SingleDeviceStrategy()
>>> trainer = Trainer(strategy=strategy)

property devices: List[Any]

List of devices used by this strategy.

property name: str

Strategy name.

property num_devices: int

Number of devices used by this strategy.

setup(model, optimizer)

Set up the model and optimizer for distributed training.

Parameters:

• model (Any) – The model to distribute.

• optimizer (Any) – The optimizer to distribute.

Returns:

The distributed model and optimizer.

Return type:

Tuple[Any, Any]

training_step(model, optimizer, batch, loss_fn, param_states)

Single device training step using brainstate transforms.

Return type:

Tuple[Any, PyTree]