[Core] Delay Loading NGram-Blocking on GPU

parlai/core/torch_generator_agent.py

@@ -1440,11 +1440,10 @@ def __init__(
         self.eos_top = False
         self.eos_top_ts = None
         self.n_best_counter = 0
-        self.gpu_beam_blocking = gpu_beam_blocking
+        self.gpu_beam_blocking = gpu_beam_blocking and torch.cuda.is_available()
         self.partial_hyps = torch.tensor([[self.bos] for i in range(beam_size)])
         if self.gpu_beam_blocking:
             self.partial_hyps = self.partial_hyps.cuda()
-        if torch.cuda.is_available():
             self.no_repeat_ngram_op = NGramRepeatBlock()
 
     def set_context(self: TSType, context: torch.LongTensor) -> TSType:

parlai/ops/ngram_repeat_block.py

@@ -16,26 +16,6 @@
 import os
 from torch.utils.cpp_extension import load
 
-current = os.getcwd()
-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=[
-            '../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)
-
 
 class NGramRepeatBlock(nn.Module):
     """
@@ -44,6 +24,24 @@ class NGramRepeatBlock(nn.Module):
 
     def __init__(self):
         super(NGramRepeatBlock, self).__init__()
+        current = os.getcwd()
+        abspath = os.path.abspath(__file__)
+        dname = os.path.dirname(abspath)
+        os.chdir(dname)
+
+        try:
+            self.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}")
+            self.ngram_repeat_block_cuda = None
+
+        os.chdir(current)
 
     def forward(
         self,
@@ -75,7 +73,7 @@ def forward(
         hypothesis = hypothesis.contiguous()
         context = context.contiguous()
         lprobs = lprobs.contiguous()
-        return ngram_repeat_block_cuda.forward(
+        return self.ngram_repeat_block_cuda.forward(
             hypothesis,
             context,
             lprobs,