Update utils.py

utils.py

@@ -277,7 +277,8 @@ class MolecularGenerationModel():
             del batch['property_names']
             del batch['non_normalized_properties']
 
-            batch = {k: v.to(self.model.device) for k, v in batch.items()}
+            batch['input_ids'] = batch['input_ids'].to(self.model.device)
+            #batch = {k: v.to(self.model.device) for k, v in batch.items()}
 
             input_length = batch['input_ids'].shape[1]
             steps = 1024 - input_length