deepmd.pd.utils.auto_batch_size

Classes

AutoBatchSize

Auto batch size.

Module Contents

class deepmd.pd.utils.auto_batch_size.AutoBatchSize(initial_batch_size: int = 1024, factor: float = 2.0, *, silent: bool = False)

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

Check if GPU is available.

Returns:

bool

True if GPU is available

is_oom_error(e: Exception) → bool

Check if the exception is an OOM error.

Parameters:

eException

Exception