mirror of
https://github.com/wassname/Castor.git
synced 2026-10-04 12:10:17 +08:00
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:
1 parent
a2904efe5a
commit
a601722768
1 file changed
+3
@@ -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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user