Update app.py

app.py

@@ -32,7 +32,7 @@ from glob import glob
 import shutil
 import pandas as pd
 import numpy as np
-import matplotlib.pyplot as plt
+#import matplotlib.pyplot as plt
 import torch.nn.functional as F
 
 # Set the device (GPU or CPU)
@@ -166,7 +166,7 @@ prod_dl = DataLoader(prod_ds, batch_size=BATCH_SIZE)
 
 # Testing the dataloader
 prod_inputs = next(iter(prod_dl))
-print(prod_inputs['pixel_values'].shape) 
+st.write(prod_inputs['pixel_values'].shape) 
 
 ## Read image from Camera
 #enable = st.checkbox("Enable camera")