Option to disable the CuDNN backend (#52)

Some of the kernels in this backend are known to be non-deterministic. It's also not clear to me whether it's enabled in a default install or not, enabled is set to True by default, but I think ultimately it's a runtime choice.
This commit is contained in:
Matt Crane authored and GitHub committed 2017-09-22 11:19:30 +09:00
1 parent a2904efe5a
commit a601722768
1 file changed
+3
+3
View File
@@ -119,6 +119,7 @@ if __name__ == "__main__":
ap.add_argument("--index-for-corpusIDF", help="fetches idf from Index. provide index path. will\
generate a vocabFile")
ap.add_argument('--seed', help='Random seed', type=int, default=1234)
ap.add_argument('--nocudnn', help='Disable the CuDNN backend' action="store_true")
args = ap.parse_args()
@@ -126,6 +127,8 @@ if __name__ == "__main__":
np.random.seed(args.seed)
if args.cuda and torch.cuda.is_available():
torch.cuda.manual_seed(args.seed)
if args.nocudnn:
torch.backends.cudnn.enabled = False
torch.set_num_threads(args.num_threads)