have one error when use 128*128 datasets about faceshifter_mine

[dayu1979]

Traceback (most recent call last):
File "train_aei.py", line 107, in
Y, Xt_attr = G(Xt, embed)
File "/home/dayu/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/dayu/D/src/FaceShifter_mine/network/aei.py", line 112, in forward
Y = self.generator(attr, z_id)
File "/home/dayu/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/dayu/D/src/FaceShifter_mine/network/aei.py", line 90, in forward
x = self.add1(x, z_att[0], z_id)
File "/home/dayu/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/dayu/D/src/FaceShifter_mine/network/add.py", line 64, in forward
x = self.add1(h, z_att, z_id)
File "/home/dayu/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/dayu/D/src/FaceShifter_mine/network/add.py", line 34, in forward
A = att_gamma * h_norm + att_beta
File "/home/dayu/anaconda3/lib/python3.7/site-packages/apex/amp/wrap.py", line 53, in wrapper
return orig_fn(*args, **kwargs)
File "/home/dayu/anaconda3/lib/python3.7/site-packages/apex/amp/wrap.py", line 53, in wrapper
return orig_fn(*args, **kwargs)
RuntimeError: The size of tensor a (4) must match the size of tensor b (2) at non-singleton dimension 3