added a device flag to predict

BrainMapperUNet.py

@@ -139,7 +139,7 @@ class BrainMapperUNet(nn.Module):
 
         return next(self.parameters()).is_cuda
         
-    def predict(self, X):
+    def predict(self, X, device= 0):
         """Post-training Output Prediction
         
         This function predicts the output of the of the U-net post-training