set cudnn to deterministic when use seed

Xiao-Chenguang · Nov 5, 2024 · ad30d8f · ad30d8f
1 parent 763e415
commit ad30d8f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/tiny_image_net_demo.py b/examples/tiny_image_net_demo.py
@@ -24,6 +24,8 @@ def test_tiny_image_net_cuda(data_path: str = ""):
 
     if args.SEED >= 0:
         torch.manual_seed(args.SEED)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
 
     assert torch.cuda.is_available(), "CUDA is not available"