ogbn-products inference OOM

[LukeLIN-web]

    @torch.no_grad()
    def inference(self, x_all, device, subgraph_loader):
        for i in range(self.num_layers):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to(device)
                x = x_all[n_id].to(device)
                x_target = x[:size[1]]
                x = self.convs[i]((x, x_target), edge_index)
                if i != self.num_layers - 1:
                    x = F.relu(x)
                xs.append(x)
            x_all = torch.cat(xs, dim=0)
        return x_all

    with torch.no_grad():
        out = model.inference(x, rank, subgraph_loader).cpu()
    y = y.cpu()
    res = out.argmax(dim=-1) == y #  big graph may oom
RuntimeError: [enforce fail at alloc_cpu.cpp:73] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 5997743042841 bytes. Error code 12 (Cannot allocate memory)

I don't know what is difference with pytorch_geometric/examples/ogbn_products_sage.py Why it need 5997G memory?

[rusty1s]

Haven‘t run the code yet, but this looks like a bug to me. What‘s the shape of output and y?