Here’s my code:
    model = models.resnet50(pretrained=pretrain_status)
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(
      nn.Dropout(dropout_rate),
      nn.Linear(num_ftrs, 1))

  criterion = nn.BCELossWithLogits() 
  optimizer = torch.optim.SGD(model.parameters(), lr =lr)

  model.train()
  for image_batch, label_batch in tqdm(train_dl):
    
    optimizer.zero_grad()
    image_batch, label_batch = image_batch.to(DEVICE), label_batch.to(DEVICE)
    pred_batch = model(image_batch)
   
    label_batch = label_batch.view(-1, 1)

    batch_loss = criterion(pred_batch.double(), label_batch.double())
    total_train_loss += batch_loss.item()

    batch_loss.backward()
    optimizer.step()

    pred_batch[pred_batch >= 0] = 1
    pred_batch[pred_batch < 0] = 0

I know there’s a mistake in there but I can’t seem to find it