expected input[8, 77, 80, 80] to have 128 channels, but got 77 channels instead

[naivewx]

when I use 'Torch-Pruning/examples/coco_yolo/train_with_pruning.py' to pure yolov5s , will have error'RuntimeError: Given groups=1, weight of size [255, 128, 1, 1], expected input[8, 77, 80, 80] to have 128 channels, but got 77 channels instead' @VainF @JonnyKong @Serjio42 ,could you help check this issue?

[YellenYeason]

put all the files which include Torch-Pruning to your file which include yolov5/train.py.