add device string

streamlit/pages/chest.py

@@ -16,15 +16,15 @@ guidance_scale = st.number_input("Guidance scale", min_value=1, max_value=10, va
 seed = st.number_input("Seed", min_value=0, max_value=None, value=1)
 cond_str = st.radio("Cardiomegaly", ('Yes', 'No'), index=1, help="Conditioned on 'cardiomegaly' or 'no cardiomegaly'", horizontal=True)
 torch.manual_seed(seed)
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = torch.device(device_str)
 
 @st.cache(allow_output_mutation = True)  
 def init_pipeline():
     pipeline = DiffusionPipeline.load_from_checkpoint('runs/chest_diffusion/last.ckpt')
     return pipeline
 
-if st.button('Sample'):
+if st.button(f'Sample (using {device_str})'):
     cond = {'Yes':1, 'No':0}[cond_str]
     condition = torch.tensor([cond]*n_samples, device=device)
     un_cond = torch.tensor([1-cond]*n_samples, device=device)

streamlit/pages/eye.py

@@ -17,14 +17,15 @@ guidance_scale = st.number_input("Guidance scale", min_value=1, max_value=10, va
 seed = st.number_input("Seed", min_value=0, max_value=None, value=1)
 cond_str = st.radio("Glaucoma", ('Yes', 'No'), index=1, help="Conditioned on 'referable glaucoma' or 'no referable glaucoma'", horizontal=True)
 torch.manual_seed(seed)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = torch.device(device_str)
 
 @st.cache(allow_output_mutation = True)  
 def init_pipeline():
     pipeline = DiffusionPipeline.load_from_checkpoint('runs/eye_diffusion/last.ckpt')
     return pipeline
 
-if st.button('Sample'):
+if st.button(f'Sample (using {device_str})'):
     cond = {'Yes':1, 'No':0}[cond_str]
     condition = torch.tensor([cond]*n_samples, device=device)
     un_cond = torch.tensor([1-cond]*n_samples, device=device)