a quick fix to wholememory tensor gather default data type

rapidsai · May 27, 2024 · 890943f · 890943f
1 parent 7352f1c
commit 890943f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/pylibwholegraph/pylibwholegraph/torch/tensor.py b/python/pylibwholegraph/pylibwholegraph/torch/tensor.py
@@ -67,7 +67,7 @@ def gather(self,
         embedding_count = indice.shape[0]
         current_cuda_device = "cuda:%d" % (torch.cuda.current_device(),)
         output_dtype = (
-            force_dtype if force_dtype is not None else self.embedding_tensor.dtype
+            force_dtype if force_dtype is not None else self.dtype
         )
         output_tensor = torch.empty(
             [embedding_count, embedding_dim],