Train error on Custom Dataset

[Kiều Trọng Thiện]

Hi, thank you for your project. I'm trying to train on my dataset but it have this error, i don't know why and i trying to find a reason to fix it.

RuntimeError: mat1 dim 1 must match mat2 dim 0
ret = torch.addmm(bias, input, weight.t())
x_warp, last_flow, last_flow_all, flow_all, delta_list, x_all, x_edge_all, delta_x_all, delta_y_all = self.aflow_net(image_input, image_edge, image_pyramids, cond_pyramids)
RuntimeError File "/home/qconnection/miniconda3/envs/ohrey/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl

And i tracked to this code:
x_warp, last_flow, last_flow_all, flow_all, delta_list, x_all, x_edge_all, delta_x_all, delta_y_all = self.aflow_net(image_input, image_edge, image_pyramids, cond_pyramids)

My shape size:

torch.Size([4, 45, 512, 384]) -> concat.cuda()
torch.Size([4, 3, 512, 384]) -> clothes.cuda()
torch.Size([4, 1, 512, 384]) -> pre_clothes_edge.cuda()

My dataset trained on PFAFN and everything is work fine, don't need to change anything in PFAFN source code. But i bring it to here and it show error