fix feature padding in GCPN

Author: KiddoZhu

torchdrug/tasks/generation.py

@@ -1269,11 +1269,11 @@ def _apply_action(self, graph, off_policy, max_resample=10, verbose=0, min_node=
             data_dict.pop(key)
         # pad 0 for node / edge attributes
         for k, v in data_dict.items():
-            if meta_dict[k] == "node":
+            if "node" in meta_dict[k]:
                 shape = (len(new_atom_type), *v.shape[1:])
                 new_data = torch.zeros(shape, dtype=v.dtype, device=self.device)
                 data_dict[k] = functional._extend(v, graph.num_nodes, new_data, has_new_node)[0]
-            if meta_dict[k] == "edge":
+            if "edge" in meta_dict[k]:
                 shape = (len(new_edge_list) * 2, *v.shape[1:])
                 new_data = torch.zeros(shape, dtype=v.dtype, device=self.device)
                 data_dict[k] = functional._extend(v, graph.num_edges, new_data, has_new_edge * 2)[0]