Is this page helpful?

`core.optimizer.optimizer_cuda_graph`#

CUDA graph wrapper for ADAM optimizer.

Module Contents#

Classes#

OptimizerCudaGraphWrapper

Wrapper class to enable FullIterationCUDAgraph.

Data#

API#

core.optimizer.optimizer_cuda_graph.logger#: ‘getLogger(…)’

class core.optimizer.optimizer_cuda_graph.OptimizerCudaGraphWrapper( optimizer_step_func, cuda_graph_warmup_steps=1, )#

Wrapper class to enable FullIterationCUDAgraph.

Initialization

curr_iteration#: 0

cuda_graph#: None

result#: None

__call__(*args, **kwargs)#

curr_iter()#: Return current training iteration.

next_iter()#: Increment current training iteration.

__del__()#