dgl.nn.NodeEmbedding.all_set_embedding not work in cpu device
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
dgl.nn.NodeEmbedding.all_set_embedding not work in cpu device.
## To Reproduce
```python
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:
```bash
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.
https://github.com/dmlc/dgl/blob/3d16000b4170fa741ed9e9667f22ba84d3493026/python/dgl/nn/pytorch/sparse_emb.py#L343-L347
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.