Repository: pyg-team/pytorch_geometric

Issue:
test-profil-it-cuda fails w/ latest
### 🐛 Describe the bug

```
_____________________________ test_profileit_cuda ______________________________
get_dataset = functools.partial(<function load_dataset at 0xffff974227a0>, '/tmp/pyg_test_datasets')
    @onlyCUDA
    @onlyOnline
    @withPackage('pytorch_memlab')
    def test_profileit_cuda(get_dataset):
        warnings.filterwarnings('ignore', '.*arguments of DataFrame.drop.*')
    
        dataset = get_dataset(name='Cora')
        data = dataset[0].cuda()
        model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
                          out_channels=dataset.num_classes).cuda()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
        @profileit('cuda')
        def train(model, x, edge_index, y):
            model.train()
            optimizer.zero_grad()
            out = model(x, edge_index)
            loss = F.cross_entropy(out, y)
            loss.backward()
            return float(loss)
    
        stats_list = []
        for epoch in range(5):
            _, stats = train(model, data.x, data.edge_index, data.y)
            assert stats.time > 0
            assert stats.max_allocated_gpu > 0
            assert stats.max_reserved_gpu > 0
            assert stats.max_active_gpu > 0
            assert stats.nvidia_smi_free_cuda > 0
            assert stats.nvidia_smi_used_cuda > 0
    
            if epoch >= 2:  # Warm-up
                stats_list.append(stats)
    
        stats_summary = get_stats_summary(stats_list)
        assert stats_summary.time_mean > 0
        assert stats_summary.time_std > 0
        assert stats_summary.max_allocated_gpu > 0
        assert stats_summary.max_reserved_gpu > 0
        assert stats_summary.max_active_gpu > 0
>       assert stats_summary.min_nvidia_smi_free_cuda > 0
E       AttributeError: 'GPUStatsSummary' object has no attribute 'min_nvidia_smi_free_cuda'
test/profile/test_profile.py:86: AttributeError

```

### Versions

latest stack as of 12 hours ago


Discussion:
https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/profile/profile.py#L34-L39
https://github.com/pyg-team/pytorch_geometric/commit/0e7458d21433f0191f50951dd2ff56e6aefa492d
it seems the profiler definition was updated but the tests weren't.
@DamianSzwichtenberg @rusty1s 

Resolving patch:
diff --git a/torch_geometric/profile/profile.py b/torch_geometric/profile/profile.py
--- a/torch_geometric/profile/profile.py
+++ b/torch_geometric/profile/profile.py
@@ -213,9 +213,7 @@ def get_stats_summary(
         max_reserved_gpu=max([s.max_reserved_gpu for s in stats_list]),
         max_active_gpu=max([s.max_active_gpu for s in stats_list]))
 
-    if all(isinstance(s, GPUStats) for s in stats_list):
-        return GPUStatsSummary(**kwargs)
-    else:
+    if all(isinstance(s, CUDAStats) for s in stats_list):
         return CUDAStatsSummary(
             **kwargs,
             min_nvidia_smi_free_cuda=min(
@@ -223,6 +221,8 @@ def get_stats_summary(
             max_nvidia_smi_used_cuda=max(
                 [s.nvidia_smi_used_cuda for s in stats_list]),
         )
+    else:
+        return GPUStatsSummary(**kwargs)
 
 
 ###############################################################################


Test patch: