Source code for imaginaire.utils.cudnn

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out
import torch.backends.cudnn as cudnn

from imaginaire.utils.distributed import master_only_print as print

[docs]def init_cudnn(deterministic, benchmark): r"""Initialize the cudnn module. The two things to consider is whether to use cudnn benchmark and whether to use cudnn deterministic. If cudnn benchmark is set, then the cudnn deterministic is automatically false. Args: deterministic (bool): Whether to use cudnn deterministic. benchmark (bool): Whether to use cudnn benchmark. """ cudnn.deterministic = deterministic cudnn.benchmark = benchmark print('cudnn benchmark: {}'.format(benchmark)) print('cudnn deterministic: {}'.format(deterministic))