Repository navigation
dgl.nn.NodeEmbedding.all_set_embedding not work in cpu device #7897
Description
🐛 Bug
dgl.nn.NodeEmbedding.all_set_embedding not work in cpu device.
To Reproduce
import torch
from dgl.nn import NodeEmbedding
emb_num = 3
emb_size = 10
ckpt_path = "model.pth"
def save(device):
emb_layer = NodeEmbedding(emb_num, emb_size, "emb", device=device,
init_func=torch.nn.init.xavier_uniform_)
init_emb = emb_layer.all_get_embedding()
print("init:\n", init_emb)
torch.save({"emb": init_emb}, ckpt_path)
def load(device):
emb_layer = NodeEmbedding(emb_num, emb_size, "emb", device=device)
print("rank:", emb_layer._rank)
checkpoint = torch.load(ckpt_path, torch.device(device))
print("ckpt:\n", checkpoint["emb"])
emb_layer.all_set_embedding(checkpoint["emb"])
print("run all_set_embedding ...")
aget_emb = emb_layer.all_get_embedding()
print("aget:\n", aget_emb)
if __name__ == "__main__":
save("cpu")
load("cpu")output:
init:
tensor([[-0.0215, -0.2327, -0.2921, 0.3921, -0.0462, 0.2429, -0.3395, -0.3044,
-0.6371, 0.5789],
[-0.6599, 0.1928, 0.2490, -0.2308, -0.6601, 0.5161, -0.5522, 0.1596,
0.5544, 0.0445],
[ 0.4633, 0.2364, 0.2244, -0.1086, -0.6691, 0.2808, -0.0796, 0.0587,
0.4433, -0.3488]])
rank: -1
ckpt:
tensor([[-0.0215, -0.2327, -0.2921, 0.3921, -0.0462, 0.2429, -0.3395, -0.3044,
-0.6371, 0.5789],
[-0.6599, 0.1928, 0.2490, -0.2308, -0.6601, 0.5161, -0.5522, 0.1596,
0.5544, 0.0445],
[ 0.4633, 0.2364, 0.2244, -0.1086, -0.6691, 0.2808, -0.0796, 0.0587,
0.4433, -0.3488]])
run all_set_embedding ...
aget:
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])Environment
- DGL Version (e.g., 1.0): 2.4.0
- Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): 2.3.1
- OS (e.g., Linux): Linux
- How you installed DGL (
conda,pip, source): pip - Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version (if applicable):
- GPU models and configuration (e.g. V100):
- Any other relevant information:
Additional context
There is no code implementation for the rank=-1 case.
dgl/python/dgl/nn/pytorch/sparse_emb.py
Lines 343 to 347 in 3d16000
Activity
github-actions commentedon Sep 18, 2025
This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you