protect load
"""
Wrapper for ngram_repeat_block cuda extension.
"""
+import parlai.utils.logging as logging
import torch
from torch import nn
@@ -20,6 +21,8 @@ abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)
+
+try:
ngram_repeat_block_cuda = load(
name='ngram_repeat_block_cuda',
sources=[
@@ -27,6 +30,10 @@ ngram_repeat_block_cuda = load(
'../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)
