Unable to load the trained model for visualization

[amayuelas]

When running the visualize script, the system does not seem to be loading the model correctly

python script/visualize.py -c config/fb15k237_visualize.yaml --checkpoint path/to/epoch/model_epoch_10.pth

This is what it returns:

RuntimeError: Error(s) in loading state_dict for LogicalQuery:
	While copying the parameter named "fact_graph", expected torch.Tensor or Tensor-like object from checkpoint but received <class 'torchdrug.data.graph.Graph'>
	While copying the parameter named "graph", expected torch.Tensor or Tensor-like object from checkpoint but received <class 'torchdrug.data.graph.Graph'>