Shape error in target

I was attempting to take a shot at the SETI ET 2021 competition in Kaggle now, as a practice of course. However, I am running into a shape error which I don’t seem to understand.

The notebook is here for reference - SETI ET 2023 w/ PyTorch | Kaggle

nn.BCEWithLogitsLoss expects model outputs and targets in the same shape while your model output is 4-dimensional. In case you are working on a binary classification use case make sure the model output has the shape [batch_size, 1].

Thanks for the response.

It is a binary classification problem. I have tried a few different variations but not able to figure out what is causing this shape error. I am new to PyTorch and hence having a hard time debugging it (I am missing something obvious but not able to figure it out)

I am using the swintranformer from timm. Is there a way where I can check the shapes layer by layer and squish it to the required shape at the end?

Reference - SETI ET 2023 w/ PyTorch | Kaggle