Shape error

[showaker]

Hello, thanks for your great work!
I came up with this problem while run "demo.py":

before pool fea: torch.Size([1, 512, 17, 50])
torch.Size([1, 8, 17, 50])
Traceback (most recent call last):
File "demo.py", line 116, in
pred = net(imgs)
File "/root/miniconda3/envs/lane-det/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/root/autodl-tmp/Ultra-Fast-Lane-Detection-V2/model/model_culane.py", line 58, in forward
fea = fea.view(-1, self.input_dim)
RuntimeError: shape '[-1, 4000]' is invalid for input of size 6800

I modified "demo.py" for predict imgs, but maybe it isn't the question. Could u give me some advice? Here is my modified "demo.py"

if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True

    args, cfg = merge_config()
    cfg.batch_size = 1
    print('setting batch_size to 1 for demo generation')

    dist_print('start testing...')
    assert cfg.backbone in ['18','34','50','101','152','50next','101next','50wide','101wide']

    if cfg.dataset == 'CULane':
        cls_num_per_lane = 18
    elif cfg.dataset == 'Tusimple':
        cls_num_per_lane = 56
    else:
        raise NotImplementedError
    net = get_model(cfg)
    state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
    compatible_state_dict = {}
    for k, v in state_dict.items():
        if 'module.' in k:
            compatible_state_dict[k[7:]] = v
        else:
            compatible_state_dict[k] = v

    net.load_state_dict(compatible_state_dict, strict=False)
    net.eval()

    img_transforms = transforms.Compose([
        transforms.Resize((int(cfg.train_height / cfg.crop_ratio), cfg.train_width)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    # 只考虑CULane数据集 
    if cfg.dataset == 'CULane':
        img_w, img_h = 1640, 590
    else:
        raise NotImplementedError
    single_image_path = './test_pic/00420.jpg' # culane自己的数据集
    vis = cv2.imread(single_image_path)
    # vis = cv2.resize(vis,(1640, 590))
    vis_pil = Image.fromarray(cv2.cvtColor(vis, cv2.COLOR_BGR2RGB))
    imgs = img_transforms(vis_pil).unsqueeze(0).cuda()  # 确保图像是批处理形式，并移动到CUDA设备
    
    start_time = time.time()

    with torch.no_grad():
        pred = net(imgs)
    
    end_time_model = time.time()
    
    coords = pred2coords(pred, cfg.row_anchor, cfg.col_anchor, original_image_width = img_w, original_image_height = img_h)
    for lane in coords:
        for coord in lane:
            cv2.circle(vis, coord, 5, (0, 255, 0), -1)
    
    output_name = "out"  # 文件名前缀
    file_extension = ".jpg"  # 文件扩展名
    output_folder = "./out_pic/"  # 输出文件夹路径
    file_index = 1  # 文件序号或索引
    output_file_path = f"{output_folder}{output_name}_{single_image_path.split('/')[-1]}"
    cv2.imwrite(output_file_path, vis)

Thanks a lot!