Data type and device transfers

metatrain.utils.transfer.batch_to(systems: List[System], targets: Dict[str, TensorMap], extra_data: Dict[str, TensorMap] | None = None, dtype: dtype | None = None, device: device | None = None) → Tuple[List[System], Dict[str, TensorMap], Dict[str, TensorMap] | None][source]

Changes the systems and targets to the specified floating point data type.

Parameters:

  • systems (List[System]) – List of systems.

  • targets (Dict[str, TensorMap]) – Dictionary of targets.

  • extra_data (Dict[str, TensorMap] | None) – Optional dictionary of extra data.

  • dtype (dtype | None) – Desired floating point data type.

  • device (device | None) – Device to move the data to.

Returns:

The systems, targets, and extra data with moved to the specified device and with the desired data type.

Return type:

Tuple[List[System], Dict[str, TensorMap], Dict[str, TensorMap] | None]