to() got an unexpected keyword argument 'non_blocking' for DGLGraph
🐛 Bug

To Reproduce
I use dgl library to make a gnn and batch the DGLGraph.
No problem during training, but in test, I got a TypeError: to() got an unexpected keyword argument 'non_blocking'
<class 'dgl.graph.DGLGraph'> .to() function has no keyword argument 'non_blocking'

Code sample

Expected behavior

Environment

OS: Linux
CUDA: 10.1
Python Version: 3.7
PyTorch Version: 1.5.1
DGL Version: 0.4.3post2
PyTorch-Lightning Version: 0.8.5

Additional context
   File "../src/main.py", line 131, in <module>
    run(params)
  File "../src/main.py", line 92, in run
    trainer.test(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in test
    results = self.__test_given_model(model, test_dataloaders)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1346, in __test_given_model
    results = self.fit(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1003, in fit
    results = self.single_gpu_train(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 186, in single_gpu_train
    results = self.run_pretrain_routine(model)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in run_pretrain_routine
    results = self.run_evaluation(test_mode=True)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 391, in run_evaluation
    eval_results = self._evaluate(self.model, dataloaders, max_batches, test_mode)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 293, in _evaluate
    output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 458, in evaluation_forward
    batch = self.transfer_batch_to_gpu(batch, root_gpu)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 159, in transfer_batch_to_gpu
    return self.__transfer_batch_to_device(batch, device)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 164, in __transfer_batch_to_device
    return model.transfer_batch_to_device(batch, device)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/core/hooks.py", line 242, in transfer_batch_to_device
    return move_data_to_device(batch, device)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 109, in move_data_to_device
    return apply_to_collection(batch, dtype=(TransferableDataType, Batch), function=batch_to)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 40, in apply_to_collection
    for k, v in data.items()})
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 40, in <dictcomp>
    for k, v in data.items()})
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 35, in apply_to_collection
    return function(data, *args, **kwargs)
  File "/home/jiangyize/miniconda3/envs/galixir/lib/python3.7/site-packages/pytorch_lightning/utilities/apply_func.py", line 107, in batch_to
    return data.to(device, non_blocking=True)
TypeError: to() got an unexpected keyword argument 'non_blocking'