Update UNet.py

UNet.py

@@ -1,3 +1,15 @@
+import torch # ; print('Using torch version -', torch.__version__)
+if torch.cuda.is_available():
+    device = 'cuda'
+else:
+    device = 'cpu'
+from torch.nn import Module, Conv2d, ReLU, ModuleList, MaxPool2d, ConvTranspose2d, BCELoss, BCEWithLogitsLoss, functional as F
+from torch.optim import Adam
+from torchvision import transforms
+from torchvision.transforms import CenterCrop
+from torch.utils.data import Dataset, DataLoader
+import cv2
+
 height, width = 256, 256
 class Block(Module):
     def __init__(self, inChannels, outChannels):