deepmd.dpmodel.utils.batch_size

deepmd.dpmodel.utils.batch_size#

Classes#

AutoBatchSize

Automatic batch size for NumPy.

Module Contents#

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

Bases: deepmd.utils.batch_size.AutoBatchSize

Automatic batch size for NumPy.

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

Contents