fix the bug in inference.py

inference.py

@@ -7,6 +7,7 @@ import torch
 from torch.nn.utils.rnn import pad_sequence
 
 from fairseq import checkpoint_utils, options, tasks, utils
+from eet.fairseq.transformer import EETTransformerDecoder
 
 Batch = namedtuple('Batch', 'ids src_tokens src_lengths')
 
@@ -77,7 +78,12 @@ class Inference(object):
         use_cuda = torch.cuda.is_available() and not args.cpu
         self.use_cuda = use_cuda
 
-        # Optimize ensemble for generation
+        model_path = args.path
+        checkpoint = torch.load(model_path.replace("best.pt", "best_part_1.pt"))
+        checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_2.pt")))
+        checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_3.pt")))
+        torch.save(checkpoint, model_path)
+        
         state = torch.load(args.path, map_location=torch.device("cpu"))
         cfg_args = eval(str(state["cfg"]))["model"]
         del cfg_args["_name"]
@@ -97,7 +103,6 @@ class Inference(object):
                       "max_batch":eet_batch_size,      
                       "full_seq_len":eet_seq_len}
         print(model_args)
-        from eet.fairseq.transformer import EETTransformerDecoder
         eet_model = EETTransformerDecoder.from_fairseq_pretrained(model_id_or_path = args.path,
                                                     dictionary = self.src_dict,args=model_args, 
                                                     config = eet_config,