Spaces:

ChemFM
/

molecular_conditional_generation

Running on Zero

feiyang-cai commited on Jan 25

Commit

6967f07

1 Parent(s): 21ff76d

debug

Files changed (1) hide show

utils.py CHANGED Viewed

@@ -277,6 +277,8 @@ class MolecularGenerationModel():
             del batch['property_names']
             del batch['non_normalized_properties']
             input_length = batch['input_ids'].shape[1]
             steps = 1024 - input_length

             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