Cast to float to ensure that it works when the frame is empty because…

… the default dtype is object
rapidsai · Feb 5, 2024 · e4a5aa5 · e4a5aa5
1 parent b04c7cf
commit e4a5aa5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/python/cugraph/cugraph/gnn/dgl_extensions/dgl_uniform_sampler.py b/python/cugraph/cugraph/gnn/dgl_extensions/dgl_uniform_sampler.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-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
@@ -120,9 +120,9 @@ def sample_neighbors(
             return self._get_edgeid_type_d(sampled_df)
         else:
             return (
-                sampled_df[src_n].values,
-                sampled_df[dst_n].values,
-                sampled_df["indices"].values,
+                sampled_df[src_n].astype("float").values,
+                sampled_df[dst_n].astype("float").values,
+                sampled_df["indices"].astype("float").values,
             )
 
     def _get_edgeid_type_d(self, df):