My suggestion  is , Can it be more optimized if we do computations on tensors instead of arrays like in this line here running_test_acc += (y_pred.argmax(dim=1) == y_batch).sum().item(). So, what’s your take and if I’m correct I want the computation to be done in tensors. Here is the code:def train_model(model, train_loader, test_loader, epochs = None, learning_rate = None, optimization_technique = None,  patience=None, scheduler_patience=None,num_batches = None, **kwargs):
    model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
   
    if optimization_technique == 'learning_rate_scheduler' and scheduler_patience:
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=scheduler_patience, verbose=True)
    else:
      scheduler = None 

    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []

    best_loss = float('inf')
    stopping_counter = 0
    
    # Enable cudnn benchmark
    torch.backends.cudnn.benchmark = True

    with_dropout = model.classifier[6]
    without_dropout = nn.Sequential(*[layer for layer in model.classifier if layer != with_dropout])

    for epoch in range(epochs):
        running_train_loss = 0.0
        running_train_acc = 0
        num_batches_train = 0
        
        #for X_batch, y_batch in itertools.islice(train_loader, 0, num_batches):
        for X_batch, y_batch in train_loader:
            
            # Move batch to device
            X_batch, y_batch = X_batch.to(device, non_blocking=True), y_batch.to(device, non_blocking=True)

            optimizer.zero_grad()
            y_pred = model(X_batch.half()) #if kwargs.get('use_mixed_precision', False) else X_batch)
            #y_pred = model(X_batch)
            loss = criterion(y_pred, y_batch)
            loss.backward()
            optimizer.step()

            running_train_loss += loss.item()
            #running_train_acc += accuracy_score(y_batch.cpu().numpy(), y_pred.argmax(dim=1).cpu().numpy())
            running_train_acc += (y_pred.argmax(dim=1) == y_batch).sum().item()
            num_batches_train += 1

        train_losses.append(running_train_loss / num_batches_train)
        train_accuracies.append(running_train_acc / len(train_loader.dataset))
        print(num_batches_train, len(train_loader.dataset))

        


        # Testing segment
        running_test_loss = 0.0
        running_test_acc = 0
        num_batches_test = 0
        

        with torch.no_grad():
            for X_batch, y_batch in test_loader:
            #for X_batch, y_batch in itertools.islice(test_loader, 0, num_batches):
                
                # Move batch to device
                #X_batch, y_batch = X_batch.to(device), y_batch.to(device)        
                X_batch, y_batch = X_batch.to(device, non_blocking=True), y_batch.to(device, non_blocking=True)
                y_pred = model(X_batch.half()) #if kwargs.get(‘use_mixed_precision’, False) else X_batch)
                #y_pred = model(X_batch)
                loss = criterion(y_pred, y_batch)

                running_test_loss += loss.item()
                running_test_acc += (y_pred.argmax(dim=1) == y_batch).sum().item()
                #running_test_acc += accuracy_score(y_batch.cpu().numpy(), y_pred.argmax(dim=1).cpu().numpy())
                num_batches_test += 1

        test_losses.append(running_test_loss / num_batches_test)
        test_accuracies.append(running_test_acc / len(test_loader.dataset))

        # Early stopping
        if optimization_technique == 'early_stopping' and patience:
            if test_losses[-1] < best_loss:
                best_loss = test_losses[-1]
                stopping_counter = 0
            else:
                stopping_counter += 1
            if stopping_counter > patience:
                print(f"Early stopping at epoch {epoch+1}/{epochs}")
                break

        # Learning rate scheduler
        if optimization_technique == 'learning_rate_scheduler' and scheduler_patience and scheduler:
            scheduler.step(test_losses[-1])

        print(f"Epoch: {epoch+1}/{epochs}, Training Loss: {train_losses[-1]}, Test Loss: {test_losses[-1]}, Training Accuracy: {train_accuracies[-1]}, Test Accuracy: {test_accuracies[-1]}")

    return train_losses, train_accuracies, test_losses, test_accuracies