Initial commit

SetUp.py

@@ -0,0 +1,14 @@
+##
+import torch
+from torchvision import datasets, transforms, models
+
+def get_device():
+
+    if torch.cuda.is_available():
+        # clear cach to boost memory
+        # for new round
+        torch.cuda.empty_cache()
+        return torch.device("cuda")
+    else:
+        return torch.device("cpu")
+