Updated app.py

app.py

@@ -4,6 +4,7 @@ import os
 import shutil
 import os
 from huggingface_hub import hf_hub_download
+import torch
 
 
 # Define paths
@@ -70,6 +71,14 @@ interface = gr.Interface(
     description="Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of MS lesions."
 )
 
+# Force GPU initialization
+if torch.cuda.is_available():
+    print("CUDA is available. Initializing GPU...")
+    device = torch.device("cuda:0")
+    torch.tensor([1.0]).to(device)  # Trigger GPU initialization
+else:
+    print("No GPU available. Falling back to CPU.")
+
 # Download model files before launching the app
 download_model()