# Figure and axes object
fig, axs = plt.subplots(nrows=2, ncols=5, figsize=(20,7), sharey=True, sharex=True)

count_OK = 0
count_nOK = 0

# Loop over loader batches
for (batch_data, batch_lbls) in train_loader:

    # Loop over batch_lbls
    for i, lbl in enumerate(batch_lbls):

        # If label is 0 (nOK) plot image in row 1
        if (lbl.item() == 0) and (count_nOK < 5):
            axs[1, count_nOK].imshow(batch_data[i][0], cmap='gray')
            axs[1, count_nOK].set_title(f"nOK Part#: {str(count_nOK)}", fontsize=14)
            count_nOK += 1

        # If label is 1 (OK) plot image in row 0
        elif (lbl.item() == 1) and (count_OK < 5):
            axs[0, count_OK].imshow(batch_data[i][0], cmap='gray')
            axs[0, count_OK].set_title(f"OK Part#: {str(count_OK)}", fontsize=14)
            count_OK += 1

    # If both counters are >=5 stop looping
    if (count_OK >=5) and (count_nOK >=5):
        break

# Config the plot canvas
fig.suptitle("Sample plot of OK and nonOK Parts", fontsize=24)
plt.setp(axs, xticks=[], yticks=[]) 
plt.show()