Better support for torch.no_grad (since volatile is deprecated)

fairseq/modules/linearized_convolution.py

@@ -69,6 +69,7 @@ class LinearizedConvolution(ConvTBC):
             # append next input
             self.input_buffer[:, -1, :] = input[:, -1, :]
             input = utils.volatile_variable(self.input_buffer)
+        with utils.maybe_no_grad():
             output = F.linear(input.view(bsz, -1), weight, self.bias)
         return output.view(bsz, 1, -1)
 

fairseq/sequence_generator.py

@@ -71,6 +71,7 @@ class SequenceGenerator(object):
             srclen = input['src_tokens'].size(1)
             if timer is not None:
                 timer.start()
+            with utils.maybe_no_grad():
                 hypos = self.generate(input['src_tokens'], beam_size=beam_size,
                                       maxlen=int(maxlen_a*srclen + maxlen_b))
             if timer is not None:
@@ -327,6 +328,7 @@ class SequenceGenerator(object):
         avg_probs = None
         avg_attn = None
         for model, encoder_out in zip(self.models, encoder_outs):
+            with utils.maybe_no_grad():
                 decoder_out, attn = model.decoder(tokens, encoder_out)
             probs = model.get_normalized_probs(decoder_out[:, -1, :], log_probs=False).data
             if avg_probs is None:

fairseq/utils.py

@@ -176,7 +176,7 @@ def _upgrade_args(args):
     return args
 
 
-def maybe_no_grad(condition):
+def maybe_no_grad(condition=True):
     if hasattr(torch, 'no_grad') and condition:
         return torch.no_grad()
     # no-op context manager
@@ -185,8 +185,9 @@ def maybe_no_grad(condition):
 
 def volatile_variable(*args, **kwargs):
     if hasattr(torch, 'no_grad'):
-        with torch.no_grad():
+        # volatile has been deprecated, use the no_grad context manager instead
         return Variable(*args, **kwargs)
+    else:
         return Variable(*args, **kwargs, volatile=True)