Exemple #1
0
def allow_non_deterministic():
    prev_state = torch.are_deterministic_algorithms_enabled()
    try:
        torch.use_deterministic_algorithms(False)
        yield
    finally:
        torch.use_deterministic_algorithms(prev_state)
def main():
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    torch.use_deterministic_algorithms(True)
    pl.seed_everything(seed=123, workers=True)
    logging.basicConfig(level=logging.INFO)

    datamodule = clevr.ClevrWInstructionsDataModule(
        "data", batch_size=16, nhops=[0]
    )
    # most params obtained via inspection of dataset
    model = NSMLightningModule(
        input_size=45,
        n_node_properties=4,
        computation_steps=0 + 1,  # lul
        encoded_question_size=100,
        output_size=28,
        learn_rate=0.001,
        use_instruction_loss=False,
    )
    metric_to_track = "train_loss"
    trainer = pl.Trainer(
        gpus=-1 if torch.cuda.is_available() else 0,
        max_epochs=1000,
        callbacks=[
            pl.callbacks.EarlyStopping(monitor=metric_to_track, patience=300),
            pl.callbacks.ModelCheckpoint(save_top_k=-1,every_n_epochs=50),
        ],
    )
    trainer.fit(model, datamodule)
Exemple #3
0
def set_seed(seed=0, rank=None):
    """
    seed: basic seed
    rank: rank of the current process, using which to mutate basic seed to have a unique seed per process
    """
    os.environ['DEEPLIIF_SEED'] = str(seed)

    if seed is not None:
        if rank is not None:
            seed_final = seed + int(rank)
        else:
            seed_final = seed

        os.environ['PYTHONHASHSEED'] = str(seed_final)
        random.seed(seed_final)
        np.random.seed(seed_final)
        torch.manual_seed(seed_final)
        torch.cuda.manual_seed(seed_final)
        torch.cuda.manual_seed_all(seed_final)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        torch.use_deterministic_algorithms(True)
        print(f'deterministic training, seed set to {seed_final}')
    else:
        print(f'not using deterministic training')
Exemple #4
0
def random_seed(seed, backends=False, deterministic_torch=True):
    """Set random seed.

    Set random seed to ``random``, ``np.random``, ``torch.backends.cudnn`` and ``torch.manual_seed``.
    Also advise torch to use deterministic algorithms.

    References:
        https://pytorch.org/docs/stable/notes/randomness.html

    Args:
        seed: Random seed.
        backends: Whether to also adapt backends. If set True cuDNN's benchmark feature is disabled. This
            causes cuDNN to deterministically select an algorithm, possibly at the cost of reduced performance.
            Also the selected algorithm is set to run deterministically.
        deterministic_torch: Whether to set PyTorch operations to behave deterministically.

    """
    from torch import manual_seed
    from torch.backends import cudnn
    import random
    random.seed(seed)
    manual_seed(seed)
    np.random.seed(seed)
    if backends:
        cudnn.deterministic = True
        cudnn.benchmark = False
    if deterministic_torch and 'use_deterministic_algorithms' in dir(torch):
        torch.use_deterministic_algorithms(True)
Exemple #5
0
def set_determinism(
    seed: Optional[int] = NP_MAX,
    use_deterministic_algorithms: Optional[bool] = None,
    additional_settings: Optional[Union[Sequence[Callable[[int], Any]], Callable[[int], Any]]] = None,
) -> None:
    """
    Set random seed for modules to enable or disable deterministic training.

    Args:
        seed: the random seed to use, default is np.iinfo(np.int32).max.
            It is recommended to set a large seed, i.e. a number that has a good balance
            of 0 and 1 bits. Avoid having many 0 bits in the seed.
            if set to None, will disable deterministic training.
        use_deterministic_algorithms: Set whether PyTorch operations must use "deterministic" algorithms.
        additional_settings: additional settings that need to set random seed.

    Note:

        This function will not affect the randomizable objects in :py:class:`monai.transforms.Randomizable`, which
        have independent random states. For those objects, the ``set_random_state()`` method should be used to
        ensure the deterministic behavior (alternatively, :py:class:`monai.data.DataLoader` by default sets the seeds
        according to the global random state, please see also: :py:class:`monai.data.utils.worker_init_fn` and
        :py:class:`monai.data.utils.set_rnd`).
    """
    if seed is None:
        # cast to 32 bit seed for CUDA
        seed_ = torch.default_generator.seed() % MAX_SEED
        torch.manual_seed(seed_)
    else:
        seed = int(seed) % MAX_SEED
        torch.manual_seed(seed)

    global _seed
    _seed = seed
    random.seed(seed)
    np.random.seed(seed)

    if additional_settings is not None:
        additional_settings = ensure_tuple(additional_settings)
        for func in additional_settings:
            func(seed)

    if torch.backends.flags_frozen():
        warnings.warn("PyTorch global flag support of backends is disabled, enable it to set global `cudnn` flags.")
        torch.backends.__allow_nonbracketed_mutation_flag = True

    if seed is not None:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:  # restore the original flags
        torch.backends.cudnn.deterministic = _flag_deterministic
        torch.backends.cudnn.benchmark = _flag_cudnn_benchmark
    if use_deterministic_algorithms is not None:
        if hasattr(torch, "use_deterministic_algorithms"):  # `use_deterministic_algorithms` is new in torch 1.8.0
            torch.use_deterministic_algorithms(use_deterministic_algorithms)
        elif hasattr(torch, "set_deterministic"):  # `set_deterministic` is new in torch 1.7.0
            torch.set_deterministic(use_deterministic_algorithms)  # type: ignore
        else:
            warnings.warn("use_deterministic_algorithms=True, but PyTorch version is too old to set the mode.")
Exemple #6
0
def init_seed(torch_seed=100, random_seed=100, np_random_seed=100):
    '''
    Initialize all random seeds
    '''
    torch.manual_seed(torch_seed)
    random.seed(random_seed)
    np.random.seed(np_random_seed)
    torch.use_deterministic_algorithms(True)
Exemple #7
0
 def setup_method(self):
     super().setup_method()
     self.set_up_model(
         FIXTURES_ROOT / "rc" / "bidaf" / "experiment.json",
         FIXTURES_ROOT / "rc" / "squad.json",
         seed=27,
     )
     torch.use_deterministic_algorithms(True)
Exemple #8
0
def set_deterministic(deterministic: bool = True):
    # if hasattr(torch, 'backends') and hasattr(torch.backends, 'cudnn'):
    #     torch.backends.cudnn.enabled = not deterministic
    #     torch.backends.cudnn.benchmark = not deterministic
    #     torch.backends.cudnn.deterministic = deterministic
    if hasattr(torch, 'set_deterministic'):
        torch.set_deterministic(deterministic)
    else:
        torch.use_deterministic_algorithms(deterministic)
def set_deterministic():
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True

    if torch.__version__ <= Version("1.7"):
        torch.set_deterministic(True)
    else:
        torch.use_deterministic_algorithms(True)
Exemple #10
0
    def _init_deterministic(self, deterministic: bool) -> None:
        self.deterministic = deterministic
        torch.use_deterministic_algorithms(deterministic)
        if deterministic:
            # fixing non-deterministic part of horovod
            # https://github.com/PyTorchLightning/pytorch-lightning/pull/1572/files#r420279383
            os.environ["HOROVOD_FUSION_THRESHOLD"] = "0"

            # https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
Exemple #11
0
def set_seed(seed):
    """Set seeds for numpy and pytorch."""
    if seed is not None:
        if seed >= 0:
            seed_everything(seed=seed)
            torch.use_deterministic_algorithms(True)
            torch.backends.cudnn.benchmark = False
        else:
            logging.warning(
                f'the random seed should be a non-negative integer')
Exemple #12
0
def seed_prng(seed, use_cuda=False, deterministic=False):
    if deterministic:
        torch.use_deterministic_algorithms(True)
        if use_cuda:
            torch.backends.cudnn.benchmark = True
    random.seed(seed)
    numpy.random.seed(random.randint(1, 100000))
    torch.random.manual_seed(random.randint(1, 100000))
    if use_cuda is True:
        torch.cuda.manual_seed_all(random.randint(1, 100000))
 def setup_method(self):
     super().setup_method()
     self.set_up_model(
         FIXTURES_ROOT / "rc" / "dialog_qa" / "experiment.json",
         FIXTURES_ROOT / "rc" / "dialog_qa" / "quac_sample.json",
         seed=42,
     )
     self.batch = Batch(self.instances)
     self.batch.index_instances(self.vocab)
     torch.use_deterministic_algorithms(True)
Exemple #14
0
    def set_random_seed(deterministic: bool, seed: int):
        if deterministic:
            torch.use_deterministic_algorithms(True)
        # set the seed for reproducibility
        np.random.seed(seed)
        torch.manual_seed(seed)
        random.seed(seed)

        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
Exemple #15
0
def main():
    torch.use_deterministic_algorithms(True)

    torch.distributed.init_process_group("nccl")
    rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()
    torch.cuda.set_device(rank)

    explicit_nhwc = True

    dtype = torch.float16
    N, C, H, W = 1, 64, 200, 336
    Hs = ((H + 8 * world_size - 1) // (8 * world_size)) * 8
    H = Hs * world_size
    gt_bottleneck = ground_truth_bottleneck(C, dtype, explicit_nhwc)
    gt = ground_truth(N, C, H, W, dtype, 1, gt_bottleneck)

    # verify that spatial bottleneck with group_size 1 produces same results as ground truth bottleneck
    spatial_bottleneck = spatial_parallel_bottleneck(C, dtype, explicit_nhwc,
                                                     gt_bottleneck, None)
    bt = apply_to_different_bottleneck(gt, spatial_bottleneck)
    compare(gt, bt)
    #print_bottleneck_p_and_b(gt_bottleneck)
    #print_bottleneck_p_and_b(spatial_bottleneck)

    spatial_group_size = world_size
    spatial_communicator = None

    peer_pool = PeerMemoryPool(rank, world_size, spatial_group_size,
                               64 * 1024 * 1024, 2 * 1024 * 1024)

    #halex = HaloExchangerAllGather(world_size, spatial_group_size, rank, spatial_communicator)
    #halex = HaloExchangerSendRecv(world_size, spatial_group_size, rank, spatial_communicator)

    halex = HaloExchangerPeer(world_size,
                              spatial_group_size,
                              rank,
                              spatial_communicator,
                              peer_pool,
                              explicit_nhwc,
                              numSM=1)
    #print("halex.signals = %s" % (str(halex.signals)))
    # Make sure peer memory halo exchanger has finished initializing flags on all ranks before proceeding
    #torch.cuda.synchronize()
    #torch.distributed.barrier()

    bt2 = n_way_spatial(halex,
                        gt_bottleneck,
                        gt,
                        explicit_nhwc,
                        world_size,
                        rank,
                        fp32_reduce=True)
    compare(gt, bt2)
def test_equalize(device):
    torch.use_deterministic_algorithms(False)
    check_functional_vs_PIL_vs_scripted(
        F.equalize,
        F_pil.equalize,
        F_t.equalize,
        {},
        device,
        dtype=None,
        tol=1.0,
        agg_method="max",
    )
Exemple #17
0
    def _init_deterministic(self, deterministic: Union[bool, _LITERAL_WARN]) -> None:
        self.deterministic = deterministic
        if _TORCH_GREATER_EQUAL_1_11 and deterministic == "warn":
            torch.use_deterministic_algorithms(True, warn_only=True)
        else:
            torch.use_deterministic_algorithms(deterministic)
        if deterministic:
            # fixing non-deterministic part of horovod
            # https://github.com/PyTorchLightning/pytorch-lightning/pull/1572/files#r420279383
            os.environ["HOROVOD_FUSION_THRESHOLD"] = "0"

            # https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
Exemple #18
0
def set_random_seed(seed, set_cuda=False):

	_rseed = torch.initial_seed() if seed is None else seed
	rpyseed(_rseed)
	torch.manual_seed(_rseed)
	if set_cuda:
		torch.cuda.manual_seed_all(_rseed)
		# Make cudnn methods deterministic according to: https://pytorch.org/docs/stable/notes/randomness.html#cudnn
		try:
			torch.use_deterministic_algorithms(use_deterministic)
		except:
			torch.backends.cudnn.deterministic = use_deterministic
		torch.backends.cudnn.benchmark = False
Exemple #19
0
def init_seeds(seed=0, deterministic=False):
    # Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html
    # cudnn seed 0 settings are slower and more reproducible, else faster and less reproducible
    import torch.backends.cudnn as cudnn

    if deterministic and check_version(torch.__version__, '1.12.0'):  # https://github.com/ultralytics/yolov5/pull/8213
        torch.use_deterministic_algorithms(True)
        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
        # os.environ['PYTHONHASHSEED'] = str(seed)

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    cudnn.benchmark, cudnn.deterministic = (False, True) if seed == 0 else (True, False)
Exemple #20
0
def set_random_seeds(seed: int = 0) -> None:
    """
    Set torch.backends.cudnn.benchmark = True to speed up training,
    but may cause some non-determinism.

    Disabling the benchmarking feature with torch.backends.cudnn.benchmark = False
    causes cuDNN to deterministically select an algorithm, possibly at the cost of
    reduced performance.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = torch.cuda.is_available()
Exemple #21
0
    def enable_reproducibility(self, seed: int = 0) -> None:
        """Limits sources of nondeterministic behavior."""
        self._seed = seed

        torch.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)

        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.benchmark = False

        # If you want to use deterministic algorithms with CUDA, then you need to set
        # the CUBLAS_WORKSPACE_CONFIG environment variable; otherwise, Torch errors.
        # See https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility.
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
def test_choice_function_fixed(trivial_choice_problem, name):
    np.random.seed(123)
    # Pytorch does not guarantee full reproducibility in different settings
    # [1]. This may become a problem in the test suite, in which case we should
    # increase the tolerance. These are only "sanity checks" on small data sets
    # anyway and the exact values do not mean much here.
    # [1] https://pytorch.org/docs/stable/notes/randomness.html
    torch.manual_seed(123)
    # Trade off performance for better reproducibility.
    torch.use_deterministic_algorithms(True)
    x, y = trivial_choice_problem
    choice_function = choice_functions[name][0]
    params, accuracies = choice_functions[name][1], choice_functions[name][2]
    learner = choice_function(**params)
    if name == GLM_CHOICE:
        learner.fit(
            x,
            y,
            vi_params={
                "n": 100,
                "method": "advi",
                "callbacks": [CheckParametersConvergence()],
            },
        )
    else:
        learner.fit(x, y)

    s_pred = learner.predict_scores(x)
    y_pred = learner.predict_for_scores(s_pred)
    y_pred_2 = learner.predict(x)
    rtol = 1e-2
    atol = 5e-2
    assert np.isclose(0.0,
                      subset_01_loss(y_pred, y_pred_2),
                      rtol=rtol,
                      atol=atol,
                      equal_nan=False)
    for key, value in accuracies.items():
        metric = choice_metrics[key]
        if metric in metrics_on_predictions:
            pred_loss = metric(y, y_pred)
        else:
            pred_loss = metric(y, s_pred)
        assert np.isclose(value,
                          pred_loss,
                          rtol=rtol,
                          atol=atol,
                          equal_nan=False)
def set_deterministic(_seed_: int = 2020, disable_uda=False):
    random.seed(_seed_)
    np.random.seed(_seed_)
    torch.manual_seed(
        _seed_
    )  # use torch.manual_seed() to seed the RNG for all devices (both CPU and CUDA)
    torch.cuda.manual_seed_all(_seed_)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    if disable_uda:
        pass
    else:
        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
        # set a debug environment variable CUBLAS_WORKSPACE_CONFIG to ":16:8" (may limit overall performance) or ":4096:8" (will increase library footprint in GPU memory by approximately 24MiB).
        torch.use_deterministic_algorithms(True)
Exemple #24
0
def test_discrete_choice_function_fixed(trivial_discrete_choice_problem, name):
    np.random.seed(123)
    # There are some caveats with pytorch reproducibility. See the comment on
    # the corresponding line of `test_choice_functions.py` for details.
    torch.manual_seed(123)
    torch.use_deterministic_algorithms(True)
    x, y = trivial_discrete_choice_problem
    choice_function = discrete_choice_functions[name][0]
    params, accuracies = (
        discrete_choice_functions[name][1],
        discrete_choice_functions[name][2],
    )
    learner = choice_function(**params)
    if name in [MNL, NLM, GEV, PCL, MLM]:
        learner.fit(
            x,
            y,
            vi_params={
                "n": 100,
                "method": "advi",
                "callbacks": [CheckParametersConvergence()],
            },
        )
    else:
        learner.fit(x, y)
    s_pred = learner.predict_scores(x)
    y_pred = learner.predict_for_scores(s_pred)
    y_pred_2 = learner.predict(x)
    rtol = 1e-2
    atol = 5e-2
    assert np.isclose(0.0,
                      subset_01_loss(y_pred, y_pred_2),
                      rtol=rtol,
                      atol=atol,
                      equal_nan=False)
    for key, value in accuracies.items():
        metric = metrics[key]
        if metric in metrics_on_predictions:
            pred_loss = metric(y, y_pred)
        else:
            pred_loss = metric(y, s_pred)
        assert np.isclose(value,
                          pred_loss,
                          rtol=rtol,
                          atol=atol,
                          equal_nan=False)
Exemple #25
0
def set_sibyl_seed(seed=41):
    import torch
    import random
    import numpy as np

    SIBYL_SEED = seed

    # torch
    torch.manual_seed(SIBYL_SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(SIBYL_SEED)
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    # python
    random.seed(SIBYL_SEED)
    #numpy
    np.random.seed(SIBYL_SEED)
Exemple #26
0
def seed_everything(seed, deterministic=False, compensation="memory") -> None:
    """
    Function that sets seed for pseudo-random number generators in:
    pytorch, numpy, python.random

    Adapted from pytorch-lightning
    https://pytorch-lightning.readthedocs.io/en/latest/_modules/pytorch_lightning/utilities/seed.html#seed_everything

    Args:
        seed (int): Value of the seed for all pseudo-random number generators
        deterministic (bool): If set to True will raise an error if non-deterministic behaviour is encountered
        compensation (str): Chooses which computational aspect is affected when deterministic is set to True.
            Must be chosen between time and memory.

    Raises:
        ClinicaDLConfigurationError: if compensation is not in {"time", "memory"}.
        RuntimeError: if a non-deterministic behaviour was encountered.

    """
    from clinicadl.utils.exceptions import ClinicaDLConfigurationError

    max_seed_value = np.iinfo(np.uint32).max
    min_seed_value = np.iinfo(np.uint32).min

    if not (min_seed_value <= seed <= max_seed_value):
        seed = random.randint(min_seed_value, max_seed_value)

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if deterministic:
        torch.backends.cudnn.benchmark = False
        if compensation == "memory":
            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
        elif compensation == "time":
            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
        else:
            raise ClinicaDLConfigurationError(
                f"The compensation for a deterministic CUDA setting "
                f"must be chosen between 'time' and 'memory'.")
        torch.use_deterministic_algorithms(True)
Exemple #27
0
def test_bincount():
    """test that bincount works in deterministic setting on GPU."""
    torch.use_deterministic_algorithms(True)

    x = torch.randint(100, size=(100, ))
    # uses custom implementation
    res1 = _bincount(x, minlength=10)

    torch.use_deterministic_algorithms(False)

    # uses torch.bincount
    res2 = _bincount(x, minlength=10)

    # explicit call to make sure, that res2 is not by accident using our manual implementation
    res3 = torch.bincount(x, minlength=10)

    # check for correctness
    assert torch.allclose(res1, res2)
    assert torch.allclose(res1, res3)
Exemple #28
0
def test_object_ranker_fixed(trivial_ranking_problem, ranker_name):
    np.random.seed(123)
    # There are some caveats with pytorch reproducibility. See the comment on
    # the corresponding line of `test_choice_functions.py` for details.
    torch.manual_seed(123)
    torch.use_deterministic_algorithms(True)
    x, y = trivial_ranking_problem
    ranker, params, (loss, acc) = object_rankers[ranker_name]
    ranker = ranker(**params)
    ranker.fit(x, y)
    pred_scores = ranker.predict_scores(x)
    pred_loss = zero_one_rank_loss_for_scores_ties_np(y, pred_scores)
    rtol = 1e-2
    atol = 1e-4
    assert np.isclose(loss, pred_loss, rtol=rtol, atol=atol, equal_nan=False)
    pred = ranker.predict_for_scores(pred_scores)
    pred_2 = ranker.predict(x)
    pred_acc = zero_one_accuracy_np(pred, pred_2)
    assert np.isclose(1.0, pred_acc, rtol=rtol, atol=atol, equal_nan=False)
    pred_acc = zero_one_accuracy_np(pred, y)
    assert np.isclose(acc, pred_acc, rtol=rtol, atol=atol, equal_nan=False)
Exemple #29
0
    def compute_opacus_grad_sample(
        self,
        x: Union[torch.Tensor, PackedSequence],
        module: nn.Module,
        batch_first=True,
        loss_reduction="mean",
    ) -> Dict[str, torch.tensor]:
        """
        Runs Opacus to compute per-sample gradients and return them for testing purposes.

        Args:
            x: The tensor in input to the ``module``
            module: The ``ModelWithLoss`` that wraps the nn.Module you want to test.
            batch_first: Whether batch size is the first dimension (as opposed to the second).
                Defaults to True.
            loss_reduction: What reduction to apply to the loss. Defaults to "mean".

        Returns:
            Dictionary mapping parameter_name -> per-sample-gradient for that parameter
        """
        torch.use_deterministic_algorithms(True)
        torch.manual_seed(0)
        np.random.seed(0)

        gs_module = GradSampleModule(clone_module(module),
                                     batch_first=batch_first,
                                     loss_reduction=loss_reduction)
        grad_sample_module = ModelWithLoss(gs_module, loss_reduction)

        grad_sample_module.zero_grad()
        loss = grad_sample_module(x)
        loss.backward()

        opacus_grad_samples = {
            name: p.grad_sample
            for name, p in
            grad_sample_module.wrapped_module._module.named_parameters()
        }

        return opacus_grad_samples
Exemple #30
0
def enable_full_determinism(seed: int):
    """
    Helper function for reproducible behavior during distributed training. See
    - https://pytorch.org/docs/stable/notes/randomness.html for pytorch
    - https://www.tensorflow.org/api_docs/python/tf/config/experimental/enable_op_determinism for tensorflow
    """
    # set seed first
    set_seed(seed)

    if is_torch_available():
        #  Enable PyTorch deterministic mode. This potentially requires either the environment
        #  variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set,
        # depending on the CUDA version, so we set them both here
        os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
        torch.use_deterministic_algorithms(True)

        # Enable CUDNN deterministic mode
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    if is_tf_available():
        tf.config.experimental.enable_op_determinism()