So i am trying to train a Variational Auto Encoder, and i have created a custom loss function to train the network, the network throws the error
RuntimeError: grad can be implicitly created only for scalar outputs

heres the Loss function
def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x, reduction='none')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    the_loss = BCE + KLD
    return the_loss

and My training Code
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3,
                             weight_decay=1e-4)

num_epochs = 150

running_loss = 0
steps = 0
print_every = 1
for epoch in range(num_epochs):
    model.train()
    for data in trainloader:
        steps += 1
        img, _ = data
        img = img.cuda()
        
        decoded, mu, logvar = model(img)
        loss = loss_function(decoded, img, mu, logvar)
        
        optimizer.zero_grad()
        loss.backward()  # <------- Error On this Line
        optimizer.step()
        running_loss += loss.item()
        if steps % print_every == 0:
            model.eval()

            with torch.no_grad():
                valid_loss = validation(model, validloader)
            
            model.train()
            
            print("Epoch: {}/{}.. ".format(epoch+1, num_epochs),
                  "Training Loss: {:.4f}.. ".format(running_loss/print_every),
                  "valid Loss: {:.4f}.. ".format(valid_loss/len(validloader)))
            running_loss = 0
    if epoch % 10 == 0:
        save_im(output, 'epoch '+str(epoch))

Image size is 64, 3, 96, 96