INNER CODE UNIT · Python

CudaGraphWrapper

NVIDIA/DeepLearningExamples · PyTorch/Recommendation/DLRM/dlrm/scripts/main.py:194

class CudaGraphWrapper:
    def __init__(self, model, train_step, parallelize,
                 zero_grad, cuda_graphs=False, warmup_steps=20):

        self.cuda_graphs = cuda_graphs
        self.warmup_iters = warmup_steps
        self.graph = None
        self.stream = None
        self.static_args = None

        self.model = model

        self._parallelize = parallelize
        self._train_step = train_step
        self._zero_grad = zero_grad

        self.loss = None
        self.step = -1

View source record →