Here is the codes in text:
Define function to train a batch of IDC images
def train(trainloader, model, criterion, optimizer, scheduler):
    total_loss = 0.0
    size = len(trainloader.dataset)
    num_batches = size // trainloader.batch_size
    model.train()
    for i, (images, labels) in enumerate(trainloader):
        print(f"Training: {i}/{num_batches}", end="\r")
        
        scheduler.step()
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images) # forward pass
        loss = criterion(outputs, labels)
        total_loss += loss.item() * images.size(0)
        loss.backward()  # backprogagation
        optimizer.step()
        
    return total_loss / size

Define function to compute the accuracy on the validation set
def validate(valloader, model, criterion):
    model.eval()
    with torch.no_grad():
        total_correct = 0
        total_loss = 0.0
        size = len(valloader.dataset)
        num_batches = size // valloader.batch_size
        for i, (images, labels) in enumerate(valloader):
            print(f"Validation: {i}/{num_batches}", end="\r")
            
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            _, preds = torch.max(outputs, 1)
            total_correct += torch.sum(preds == labels.data)
            total_loss += loss.item() * images.size(0)
            
        return total_loss / size, total_correct.double() / size

main training loop
target_size=torch.rand((48,48), requires_grad=False)
input_size=torch.rand((48,48), requires_grad=False)

def fit(model, num_epochs, trainloader, valloader):
    #criterion = WeightedBCELoss() 
    criterion = StableBCELoss()
    optimizer = Adagrad(model.parameters(), lr=lr,lr_decay=lr/num_epochs)
    scheduler = OneCycleLR(optimizer, lr_range=(lr,1.), num_steps=1000)
    print("epoch\ttrain loss\tvalid loss\taccuracy")
    for epoch in range(num_epochs):
        train_loss = train(trainloader, model, criterion, optimizer, scheduler)
        valid_loss, valid_acc = validate(valloader, model, criterion)
        print(f"{epoch}\t{train_loss:.5f}\t\t{valid_loss:.5f}\t\t{valid_acc:.3f}")