code: def test_gnn_reload(tmpdir):
      import torch
      from deepchem.models.torch_models.gnn import GNNModular
  
      model_config = {
          "gnn_type": "gin",
          "num_layers": 3,
          "emb_dim": 64,
          "task": "regression",
          "mask_edge": True,
          "model_dir": tmpdir,
          "device": torch.device("cpu"),
      }
      old_model = GNNModular(**model_config)
      old_model._ensure_built()
      old_model.save_checkpoint()
      old_model_state = old_model.model.state_dict()
  
      new_model = GNNModular(**model_config)
      new_model.restore()
      new_model_state = new_model.model.state_dict()
      for key in new_model_state.keys():
          assert torch.allclose(
              old_model_state[key], new_model_state[key]
          )
  
file_path: deepchem/deepchem/models/tests/test_gnn.py
imports: import os
  import pytest
  import numpy
  import deepchem
  from deepchem.feat.molecule_featurizers import SNAPFeaturizer
label_desc: Use image node_id 6 for calling a global function with example usage: test_gnn_reload(tmpdir) without return types