protect load (facebookresearch#4779)

CZWin32768 · Sep 9, 2022 · e217dc4 · e217dc4
1 parent 58b6977
commit e217dc4
Showing 1 changed file with 14 additions and 7 deletions.
diff --git a/parlai/ops/ngram_repeat_block.py b/parlai/ops/ngram_repeat_block.py
@@ -9,6 +9,7 @@
 """
 Wrapper for ngram_repeat_block cuda extension.
 """
+import parlai.utils.logging as logging
 import torch
 from torch import nn
 
@@ -20,13 +21,19 @@
 dname = os.path.dirname(abspath)
 os.chdir(dname)
 
-ngram_repeat_block_cuda = load(
-    name='ngram_repeat_block_cuda',
-    sources=[
-        '../clib/cuda/ngram_repeat_block_cuda.cpp',
-        '../clib/cuda/ngram_repeat_block_cuda_kernel.cu',
-    ],
-)
+
+try:
+    ngram_repeat_block_cuda = load(
+        name='ngram_repeat_block_cuda',
+        sources=[
+            '../clib/cuda/ngram_repeat_block_cuda.cpp',
+            '../clib/cuda/ngram_repeat_block_cuda_kernel.cu',
+        ],
+    )
+except Exception as e:
+    logging.warning(f"Unable to load ngram blocking on GPU: {e}")
+    ngram_repeat_block_cuda = None
+
 os.chdir(current)