Some fixes for running on CPU (#221)

Summary:
Pull Request resolved: #221

Running the benchmark with CUDA_VISIBLE_DEVICES='' brought up some issues in the code when running on CPU. These are the fixes

Reviewed By: jhcross

Differential Revision: D9996497

fbshipit-source-id: e418f0d5f9bd98b97fb149d689f3571fb4b8f14d

Author: James Reed

pytorch_translate/beam_decode.py

@@ -99,6 +99,8 @@ def generate_batched_itr(
         for sample in data_itr:
             if cuda:
                 s = utils.move_to_cuda(sample)
+            else:
+                s = sample
             input = s["net_input"]
             srclen = input["src_tokens"].size(1)
             if self.use_char_source:
@@ -608,8 +610,9 @@ def gather_probs(all_translation_tokens, all_probs):
                 # The corresponding model did not use vocab reduction if
                 # possible_translation_tokens is None.
                 mapped_probs = torch.zeros(
-                    (probs.size(0), possible_translation_tokens.size(0))
-                ).cuda()
+                    (probs.size(0), possible_translation_tokens.size(0)),
+                    device=probs.device,
+                )
 
                 mapped_probs[:, inv_ind] = probs
             if avg_probs is None: