Commit with rpn gnn tuning

jwzhanggy · Oct 22, 2024 · c142e37 · c142e37
1 parent fa0ac97
commit c142e37
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tinybig/util/utility.py b/tinybig/util/utility.py
@@ -57,6 +57,7 @@ def set_random_seed(random_seed: int = 0):
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
     torch.use_deterministic_algorithms(True)
+    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
 
 
 def check_file_existence(complete_file_path):