#2777: Add determinism to torch backends when seed set.

2024-08-05 15:08:31 +01:00
parent 7d7117e624
commit 966542c2ca
1 changed files with 4 additions and 0 deletions
--- a/src/primaite/session/environment.py
+++ b/src/primaite/session/environment.py
@@ -44,6 +44,10 @@ def set_random_seed(seed: int) -> Union[None, int]:
    # if torch not installed don't set random seed.
    if sys.modules["torch"]:
        th.manual_seed(seed)
+
+    th.backends.cudnn.deterministic = True
+    th.backends.cudnn.benchmark = False
+
    return seed