Solved by ptrblck in post #2 
              
              
                Your code should work as seen here: 
model = nn.Conv2d(3, 1, 3, 1, 1)
input = torch.randn(1, 3, 100, 100)

pred = model(input)
pred_array = pred[:, :, :50, 0]
print(pred_array.shape)
> torch.Size([1, 1, 50])
loss = F.mse_loss(pred_array, torch.rand_like(pred_array))
loss.backward()

and the error me…