model to cpu

pages/1_Disentanglement.py

@@ -133,7 +133,7 @@ with input_col_3:
 
 # ---------------------------- DISPLAY COL 2 ROW 1 ------------------------------
 
-model = torch.load('./data/model_files/pytorch_model.bin')
+model = torch.load('./data/model_files/pytorch_model.bin', map_location=torch.device('cpu'))
 original_image_vec = annotations['z_vectors'][st.session_state.image_id]
 img = generate_original_image(original_image_vec, model)
 # input_image = original_image_dict['image']