Skip to content

Commit

Permalink
Update seed (#869)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #869

In this diff add some more configs to set seed

Created from CodeHub with https://fburl.com/edit-in-codehub

Reviewed By: JKSenthil

Differential Revision: D59824341

fbshipit-source-id: 0a9ac5729fa1928f4dcd1be4cfce53a0243980e7
  • Loading branch information
Ayushi Dalmia authored and facebook-github-bot committed Jul 24, 2024
1 parent fac9d15 commit 745f5cb
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 1 deletion.
2 changes: 1 addition & 1 deletion tests/framework/test_auto_unit_gpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
3 changes: 3 additions & 0 deletions torchtnt/utils/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand All @@ -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")

0 comments on commit 745f5cb

Please sign in to comment.