Mismatch output size on last layer

username: richardatpytorch
cooked: <p>I am using UNet in hub (<a href="https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/" class="inline-onebox" rel="noopener nofollow ugc">U-Net for brain MRI | PyTorch</a>) for the depth estimation task. I have no clue what is going wrong with this error</p>
<pre><code class="lang-auto">RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of size: : [64, 1, 256, 256]
</code></pre>
<p>Some context on my case: I am using pytorch lightning as a trainer and keeping the hyper params of model as suggested on the link above.</p>

username: ptrblck
cooked: <p>Could you describe your use case, which line of code raises the error, and what the target represents?</p>

username: richardatpytorch
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61d059291af00cc25912d00e3c0d3e4232a28932.png" data-download-href="https://discuss.pytorch.org/uploads/default/61d059291af00cc25912d00e3c0d3e4232a28932" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61d059291af00cc25912d00e3c0d3e4232a28932.png" alt="image" data-base62-sha1="dXiHk4YynRvPjJXUmsG0xNo5rdE" width="690" height="268" data-dominant-color="2C2C2C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1333×519 79.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Here is the full error log. The target is a depth image with values in the range [-1, 1] in size 256x256.</p>

username: ptrblck
cooked: <p>The shape mismatch is raised in <code>nn.NLLLoss</code> so check the <a href="https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html#torch.nn.NLLLoss">docs</a> and make sure the model output and target have the expected shapes.<br>
I would guess you might need to <code>squeeze</code> <code>dim1</code> in the <code>target</code> but since you didn’t explain the use case it’s just my guess.</p>

username: richardatpytorch
cooked: <p>For my use case, I am trying to use UNet for the depth estimation task. And now I understand, my label is an image size 256x256, and my model is only outputting a single label. How could I change the last layer of UNet to output a result that matches my label?</p>

username: richardatpytorch
cooked: <p>The problem is that I apply a training procedure for classification in my regression task. I fix the loss function and everything is ok now.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="richardatpytorch" data-post="6" data-topic="178147">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/richardatpytorch/48/59675_2.png" class="avatar"> richardatpytorch:</div>
<blockquote>
<p>in my regression task.</p>
</blockquote>
</aside>
<p>Ah OK, I guess you are using <code>nn.MSELoss</code> now and it works fine?</p>

username: richardatpytorch
cooked: <p>You guess it right. I forgot to check the loss function of the training pipeline at first.</p>