enable rest of the flows for use-sparse-tensor (#5258)

* update

* changelog

* enable rest flow for use-sparse-tensor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -5,7 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.2.0] - 2022-MM-DD
 ### Added
-- Added `SparseTensor` support to inference benchmark suite ([#5242](https://github.com/pyg-team/pytorch_geometric/pull/5242))
+- Added `SparseTensor` support to inference benchmark suite ([#5242](https://github.com/pyg-team/pytorch_geometric/pull/5242), [#5258](https://github.com/pyg-team/pytorch_geometric/pull/5258))
 - Added experimental mode in inference benchmarks ([#5254](https://github.com/pyg-team/pytorch_geometric/pull/5254))
 - Added node classification example instrumented with [Weights and Biases (W&B) logging](https://wandb.com) and [W&B Sweeps](https://wandb.com/sweeps) ([#5192](https://github.com/pyg-team/pytorch_geometric/pull/5192))
 - Added experimental mode for `utils.scatter` ([#5232](https://github.com/pyg-team/pytorch_geometric/pull/5232), [#5241](https://github.com/pyg-team/pytorch_geometric/pull/5241))

benchmark/inference/hetero_sage.py

@@ -18,4 +18,7 @@ def inference(self, loader, device, progress_bar=False):
             loader = tqdm(loader, desc="Inference")
         for batch in loader:
             batch = batch.to(device)
-            self.model(batch.x_dict, batch.edge_index_dict)
+            if len(batch.adj_t_dict) > 0:
+                self.model(batch.x_dict, batch.adj_t_dict)
+            else:
+                self.model(batch.x_dict, batch.edge_index_dict)

torch_geometric/nn/models/basic_gnn.py

@@ -222,7 +222,10 @@ def inference(self, loader: NeighborLoader,
             xs: List[Tensor] = []
             for batch in loader:
                 x = x_all[batch.n_id].to(device)
-                edge_index = batch.edge_index.to(device)
+                if hasattr(batch, 'adj_t'):
+                    edge_index = batch.adj_t.to(device)
+                else:
+                    edge_index = batch.edge_index.to(device)
                 x = self.convs[i](x, edge_index)[:batch.batch_size]
                 if i == self.num_layers - 1 and self.jk_mode is None:
                     xs.append(x.cpu())