UNET + YOLOv6

[geekdreamer04]

I am trying to add Unet as a preprocessing layer before the YOLOv6 architecture and when I am trying to do this, I am facing the following error. And how do I combine a Unet architecture with YOLO architecture. Any help from you here will be appreciated.
Looking forward to your support as soon as possible. Thank you

ERROR in training steps.
ERROR in training loop or eval/save model.

Training completed in 0.000 hours.
Traceback (most recent call last):
File "tools/train.py", line 112, in
main(args)
File "tools/train.py", line 102, in main
trainer.train()
File "/workspace/YOLOv61/yolov6/core/engine.py", line 75, in train
self.train_in_loop()
File "/workspace/YOLOv61/yolov6/core/engine.py", line 88, in train_in_loop
self.train_in_steps()
File "/workspace/YOLOv61/yolov6/core/engine.py", line 104, in train_in_steps
preds = self.model(images)
File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/workspace/YOLOv61/yolov6/models/yolo.py", line 39, in forward
x = self.detect(x)
File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/workspace/YOLOv61/yolov6/models/effidehead.py", line 60, in forward
x[i] = self.stemsi
File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/workspace/YOLOv61/yolov6/layers/common.py", line 102, in forward
return self.act(self.bn(self.conv(x)))
File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 446, in forward
return self._conv_forward(input, self.weight, self.bias)
File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 442, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [256, 256, 1, 1], expected input[8, 128, 20, 20] to have 256 channels, but got 128 channels instead

[ptrblck]

The error is raised in your YOLO model in:

x = self.detect(x)
...
return self.act(self.bn(self.conv(x)))

which gets an input activation with 128 channels while 256 are expected.
Could you check the shape of the UNet output and make sure it's compatible with your YOLO model?