skjaini commited on
Commit
2851f3a
·
verified ·
1 Parent(s): 95b348b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -2,6 +2,7 @@ import streamlit as st
2
  from PIL import Image
3
  from transformers import pipeline
4
  import io
 
5
 
6
  # --- Configuration ---
7
  # Specify the model
@@ -14,7 +15,7 @@ def load_pipeline():
14
  """Loads the VQA pipeline."""
15
  try:
16
  # Explicitly set device if CUDA is available, otherwise use CPU
17
- device = 0 if any(tf.config.list_physical_devices('GPU')) else -1
18
  vqa_pipeline = pipeline("visual-question-answering", model=MODEL_NAME, device=device) # Add device
19
  return vqa_pipeline
20
  except Exception as e:
 
2
  from PIL import Image
3
  from transformers import pipeline
4
  import io
5
+ import torch # Import PyTorch
6
 
7
  # --- Configuration ---
8
  # Specify the model
 
15
  """Loads the VQA pipeline."""
16
  try:
17
  # Explicitly set device if CUDA is available, otherwise use CPU
18
+ device = 0 if torch.cuda.is_available() else -1 # Use torch.cuda
19
  vqa_pipeline = pipeline("visual-question-answering", model=MODEL_NAME, device=device) # Add device
20
  return vqa_pipeline
21
  except Exception as e: