Update app.py

app.py

@@ -18,6 +18,10 @@ image_size = 256
 n_tiles = 36
 batch_size = 8
 num_workers = 4
+
+
+# Define the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class enetv2(nn.Module):
     def __init__(self, backbone, out_dim):
         super(enetv2, self).__init__()