nvalchemi.distributed.localize

nvalchemi.distributed.localize(data)

Run to_local() over every value in a model-input dict.

Localize a whole input dict in one call so a kernel that consumes it never sees a ShardTensor.

Parameters:

data (dict[str, Any]) – A model-input dict whose values may include ShardTensors.

Returns:

A shallow copy with each value run through to_local(); non-tensor entries (configs, ints) pass through unchanged.

Return type:

dict[str, Any]