Skip to content

Commit

Permalink
fix test_dataset.py for the newly defined features_transform in the s…
Browse files Browse the repository at this point in the history
…tate dict
  • Loading branch information
gcroci2 committed Oct 23, 2023
1 parent 0fcea3a commit f8e6c57
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -1019,6 +1019,11 @@ def test_inherit_info_pretrained_model_graphdataset(self):
# in the test should be inherited from the pre-trained model
inherited_params = ["node_features", "edge_features", "features_transform", "target", "target_transform", "task", "classes", "classes_to_index"]
data = torch.load(pretrained_model, pickle_module = dill, map_location=torch.device('cpu'))
if data["features_transform"]:
for _, key in data["features_transform"].items():
if key['transform'] is None:
continue
key['transform'] = eval(key['transform'])

dataset_test_vars = vars(dataset_test)
for param in inherited_params:
Expand Down

0 comments on commit f8e6c57

Please sign in to comment.