deepmd.dpmodel.utils.batch_size

Module Contents

Classes

AutoBatchSize

Automatic batch size for NumPy.
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