debug

utils.py

@@ -277,6 +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()}
+
             input_length = batch['input_ids'].shape[1]
             steps = 1024 - input_length