[torchscript] use device type when concatenating tensors (#4554)

Co-authored-by: Maryam Daneshi <[email protected]>

parlai/torchscript/modules.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from collections import defaultdict
-from typing import List, Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import torch.jit
 from parlai.agents.bart.bart import BartAgent
@@ -218,9 +218,13 @@ def forward(self, context: str, max_len: int = 128) -> str:
         if self.is_bart:
             flattened_text_vec = torch.cat(
                 [
-                    torch.tensor([self.start_idx], dtype=torch.long),
-                    flattened_text_vec,
-                    torch.tensor([self.end_idx], dtype=torch.long),
+                    torch.tensor([self.start_idx], dtype=torch.long).to(
+                        self.get_device()
+                    ),
+                    flattened_text_vec.to(self.get_device()),
+                    torch.tensor([self.end_idx], dtype=torch.long).to(
+                        self.get_device()
+                    ),
                 ],
                 dim=0,
             )