From f82c3e7e038dd1b557d115c557fe7119eab45265 Mon Sep 17 00:00:00 2001 From: zhou9402 Date: Fri, 22 Nov 2024 06:55:34 +0800 Subject: [PATCH] [Optimization] Accelerate dgl create block by bypassing sanity check (#221) Authors: - https://github.com/zhou9402 - Brad Rees (https://github.com/BradReesWork) Approvers: - Brad Rees (https://github.com/BradReesWork) URL: https://github.com/rapidsai/wholegraph/pull/221 --- python/pylibwholegraph/pylibwholegraph/torch/gnn_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pylibwholegraph/pylibwholegraph/torch/gnn_model.py b/python/pylibwholegraph/pylibwholegraph/torch/gnn_model.py index 709060c4f..fbb365046 100644 --- a/python/pylibwholegraph/pylibwholegraph/torch/gnn_model.py +++ b/python/pylibwholegraph/pylibwholegraph/torch/gnn_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2019-2023, NVIDIA CORPORATION. +# Copyright (c) 2019-2024, NVIDIA CORPORATION. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at @@ -162,6 +162,7 @@ def create_sub_graph( ), num_src_nodes=target_gid.size(0), num_dst_nodes=target_gid_1.size(0), + node_count_check=False, ) return block elif framework_name == "cugraph":