Move cudnn deterministic in eval

pytorch · Mar 21, 2024 · b6110d3 · b6110d3
1 parent 082ac0f
commit b6110d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchbenchmark/models/pytorch_unet/__init__.py b/torchbenchmark/models/pytorch_unet/__init__.py
@@ -6,7 +6,6 @@
 from torch import optim
 from typing import Tuple
 
-torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
 from .pytorch_unet.unet import UNet
@@ -89,6 +88,7 @@ def jit_callback(self):
             self.model = torch.jit.script(self.model)
 
     def eval(self) -> Tuple[torch.Tensor]:
+        torch.backends.cudnn.deterministic = True
         self.model.eval()
         with torch.no_grad():
             with torch.cuda.amp.autocast(enabled=self.args.amp):