deepmd.jax.utils.auto_batch_size

deepmd.jax.utils.auto_batch_size#

Classes#

AutoBatchSize

Auto batch size.

Module Contents#

class deepmd.jax.utils.auto_batch_size.AutoBatchSize(initial_batch_size: int = 1024, factor: float = 2.0)[source]#

Bases: deepmd.utils.batch_size.AutoBatchSize

Auto batch size.

Parameters:
initial_batch_sizeint, default: 1024

initial batch size (number of total atoms) when DP_INFER_BATCH_SIZE is not set

factorfloat, default: 2.

increased factor

is_gpu_available() → bool[source]#

Check if GPU is available.

Returns:
bool

True if GPU is available

is_oom_error(e: Exception) → bool[source]#

Check if the exception is an OOM error.

Parameters:
eException

Exception