Chap4 (CNN_working_details.ipynb) -- Squeeze on model output

[Camaltra]

On the cell to train the model

def train_batch(x, y, model, opt, loss_fn):
    model.train()
    prediction = model(x)
    batch_loss = loss_fn(prediction, y)
    batch_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return batch_loss.item()

Prediction have the outcome of [1, 1] while we give the y_train

y_train = torch.tensor([0, 1]) 

We got a miss matching size and the code doesn't run

Error Trigger

Using a target size (torch.Size([1])) that is different to the input size (torch.Size([1, 1])) is deprecated. Please ensure they have the same size.

To fix it, either

y_train = torch.tensor([0, 1]).to(device).float().unsqueeze(1)

def train_batch(x, y, model, opt, loss_fn):
    model.train()
    prediction = model(x)
    batch_loss = loss_fn(prediction.squeeze(0), y)
    batch_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return batch_loss.item()