Add torch extra-index + log

app.py

@@ -9,27 +9,33 @@
 # TODO : add the training of the vae
 # TODO : add the description of the settings
 
-import sys
 
+
+
+import gradio as gr
+import json
 import numpy as np
+import shutil
+import sys
+import tempfile
+import torch
+
 from PIL import Image, ImageDraw
-import gradio as gr
-from app.tiff_utils import extract_physical_metadata
 from app.dl_inference import inference_fn
 from app.knn import knn, segment_image, bokeh_plot_knn, color_palette
-
-import tempfile
-import shutil
-import json
-from zipfile import ZipFile
-from datetime import datetime
-
+from app.tiff_utils import extract_physical_metadata
 from collections import namedtuple
+from datetime import datetime
+from zipfile import ZipFile
 
 block_state_entry = namedtuple(
     "block_state", ["results", "knn_results", "physical_metadata"]
 )
 
+if torch_availbale := torch.cuda.is_available():
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+print(f"Is CUDA available: {torch_availbale}")
+
 if ".." not in sys.path:
     sys.path.append("..")
 

requirements.txt

@@ -170,6 +170,7 @@ tifffile==2023.2.28
 tinycss2==1.2.1
 tomli==2.0.1
 toolz==0.12.0
+--extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.2.0
 torchvision==0.17.0
 tornado==6.2