diff --git a/tests/framework/test_auto_unit_gpu.py b/tests/framework/test_auto_unit_gpu.py index 8223508614..0eafcc06b8 100644 --- a/tests/framework/test_auto_unit_gpu.py +++ b/tests/framework/test_auto_unit_gpu.py @@ -135,8 +135,8 @@ def forward(self, x): return x # so all ranks start with same initialized weights - seed(0) device = init_from_env() + seed(0) my_module = Net() auto_unit = DummyAutoUnit( diff --git a/torchtnt/utils/env.py b/torchtnt/utils/env.py index 96845f3e5d..7fb337568c 100644 --- a/torchtnt/utils/env.py +++ b/torchtnt/utils/env.py @@ -140,6 +140,7 @@ def seed(seed: int, deterministic: Optional[Union[str, int]] = None) -> None: torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) if deterministic is not None: _log.debug(f"Setting deterministic debug mode to {deterministic}") @@ -155,3 +156,5 @@ def seed(seed: int, deterministic: Optional[Union[str, int]] = None) -> None: torch.backends.cudnn.benchmark = False # reference: https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + if deterministic_debug_mode == 2: + set_float32_precision("highest")