fix trainer device while infer; fix read_data from dict when get repr… (

#214) fix trainer device while infer; fix read_data from dict when get repr without "target" key Co-authored-by: cuiyaning <cuiyaning@dp.tech>
deepmodeling · Apr 15, 2024 · 37c6ddc · 37c6ddc
1 parent c458c0f
commit 37c6ddc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/unimol_tools/unimol_tools/data/datareader.py b/unimol_tools/unimol_tools/data/datareader.py
@@ -66,7 +66,7 @@ def read_data(self, data=None, is_train=True, **params):
                     for i in range(label.shape[1]):
                         data[target_col_prefix + str(i)] = label[:,i]
 
-            _ = data.pop('target')
+            _ = data.pop('target', None)
             data = pd.DataFrame(data).rename(columns={smiles_col: 'SMILES'})
 
         elif isinstance(data, list):

diff --git a/unimol_tools/unimol_tools/predictor.py b/unimol_tools/unimol_tools/predictor.py
@@ -87,7 +87,7 @@ def get_repr(self, data=None, return_atomic_reprs=False):
                          **self.params,
                         )
         dataset = MolDataset(datahub.data['unimol_input'])
-        self.trainer = Trainer(task='repr')
+        self.trainer = Trainer(task='repr', cuda=self.device)
         repr_output = self.trainer.inference(self.model, 
                                              return_repr=True, 
                                              return_atomic_reprs=return_atomic_reprs,