Update app.py

app.py

@@ -8,6 +8,7 @@ import math
 import torch.nn.functional as F
 from tqdm.auto import tqdm
 from torchvision import transforms
+from matplotlib.pyplot import plt
 
 from torch import nn
 img_size = 64