edits

app.py

@@ -5,6 +5,8 @@ import cv2
 import torch
 import torch.utils.data as data
 from torchvision import transforms
+from torch import nn
+import torch.nn.functional as F
 
 import matplotlib.pyplot as plt
 from matplotlib import cm