Update utils.py

utils.py

@@ -208,18 +208,19 @@ class ReactionPredictionModel():
         self.forward_model.to("cuda")
 
     @spaces.GPU(duration=20)
-    def predict(self, test_loader):
+    def predict(self, test_loader, task_type):
         predictions = []
         for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
             with torch.no_grad():
                 generation_prompts = batch['generation_prompts'][0]
                 inputs = self.tokenizer(generation_prompts, return_tensors="pt", padding=True, truncation=True)
-                inputs = {k: v.to(self.retro_model.device) for k, v in inputs.items()}
-                print(inputs)
-                print(self.forward_model.device)
-                print(self.retro_model.device)
                 del inputs['token_type_ids']
+
                 if task_type == "retrosynthesis":
+                    self.retro_model.to("cuda")
+                    inputs = {k: v.to(self.retro_model.device) for k, v in inputs.items()}
+                    print(inputs)
+                    print(self.retro_model.device)
                     outputs = self.retro_model.generate(**inputs, max_length=512, num_return_sequences=10,
                                    do_sample=False, num_beams=10,
                                    eos_token_id=self.tokenizer.eos_token_id,
@@ -228,6 +229,10 @@ class ReactionPredictionModel():
                                    length_penalty=0.0,
                                    )
                 else:
+                    self.forward_model.to("cuda")
+                    inputs = {k: v.to(self.forward_model.device) for k, v in inputs.items()}
+                    print(inputs)
+                    print(self.forward_model.device)
                     outputs = self.forward_model.generate(**inputs, max_length=512, num_return_sequences=10,
                                    do_sample=False, num_beams=10,
                                    eos_token_id=self.tokenizer.eos_token_id,
@@ -281,7 +286,7 @@ class ReactionPredictionModel():
             collate_fn=self.data_collator,
         )
 
-        rank = self.predict(test_loader)
+        rank = self.predict(test_loader, task_type)
 
         return rank