tensor size of the classification head for testing and training

[tuvshu99]

Hello
I have tried to run your model with Tusimple dataset.
I have received following error while test with pretrained model.

test: 0%| | 0/696 [00:00<?, ?it/s]Traceback (most recent call last):
File "main.py", line 71, in
main()
File "main.py", line 35, in main
runner.test()
File "/Backup/test/LVLane/lanedet/engine/runner.py", line 159, in test
for i, data in enumerate(tqdm(self.test_loader, desc=f'test')):
File "/home/essys/anaconda3/envs/laneatt/lib/python3.8/site-packages/tqdm/std.py", line 1108, in iter
for obj in iterable:
.....
RuntimeError: stack expects each tensor to be equal size, but got [4] at entry 0 and [3] at entry 3

test: 0%|

[tuvshu99]

Added following instead of base_dataset.py line #67

Initialize a new list to store the updated values

    updated_category = []
    index_in_category = 0
    # Iterate over the indices of category list
    for i in range(6):
        # Check if i is within the bounds of the category list
        if index_in_category < len(category):
            # Check if all elements in the ith column of cls_label are equal to 100
            if np.all(sample['cls_label'][:, i].numpy() == 100):
                # If all elements are equal to 100, assign 0 to the corresponding entry in category
                updated_category.append(0)
            else:
                # Otherwise, keep the original value from category
                updated_category.append(category[index_in_category])
                index_in_category = index_in_category + 1
        else:
            # If i exceeds the number of elements in category, assign 0
            updated_category.append(0)

    sample['category'] = torch.LongTensor(updated_category)