RuntimeError: 'weight' must be 2-D

[billqu01]

I try the project, run it, but error, the notify below. how to fix it?
Traceback (most recent call last):
File "D:\test\PSGAN-master\demo.py", line 72, in
main()
File "D:\test\PSGAN-master\demo.py", line 57, in main
image, face = inference.transfer(source, reference, with_face=True)
File "D:\test\PSGAN-master\psgan\inference.py", line 33, in transfer
source_input, face, crop_face = self.preprocess(source)
File "D:\test\PSGAN-master\psgan\preprocess.py", line 141, in call
mask = self.face_parse.parse(cv2.resize(np_image, (512, 512)))
File "D:\test\PSGAN-master\faceutils\mask\main.py", line 37, in parse
parsing = torch.nn.functional.embedding(parsing, self.dic)
File "D:\test\PSGAN-master\venv\lib\site-packages\torch\nn\functional.py", line 2199, in embedding
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D

[hungtooc]

Here is how to fix it, it's worked with me:
in file ./faceutils/mask/main.py line 37,
change
parsing = torch.nn.functional.embedding(parsing, self.dic)
to
parsing = torch.nn.functional.embedding(parsing, torch.unsqueeze(self.dic,1))