Do not access the GPU from multiple ML libraries at the same time

LibrePhotos · Sep 16, 2023 · 8499911 · 8499911
1 parent bff7e06
commit 8499911
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/api/batch_jobs.py b/api/batch_jobs.py
@@ -3,7 +3,6 @@
 from datetime import datetime
 
 import pytz
-import torch
 from constance import config as site_config
 from django.db.models import Q
 from django_q.tasks import AsyncTask
@@ -31,6 +30,10 @@ def create_batch_job(job_type, user):
 
 
 def batch_calculate_clip_embedding(job_id, user):
+    import torch
+
+    # Only supports CPU
+    torch.device("cpu")
     lrj = LongRunningJob.objects.get(job_id=job_id)
     lrj.started_at = datetime.now().replace(tzinfo=pytz.utc)
 

diff --git a/api/im2txt/sample.py b/api/im2txt/sample.py
@@ -18,10 +18,6 @@
 decoder_path = os.path.join(im2txt_models_path, "models", "decoder-10-1000.ckpt")
 vocab_path = os.path.join(im2txt_models_path, "data", "vocab.pkl")
 
-# Device configuration
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# device = torch.device('cpu')
-
 # Load vocabulary wrapper
 with open(vocab_path, "rb") as f:
     vocab = pickle.load(f)
@@ -38,6 +34,8 @@ def load_image(image_path, transform=None):
 
 
 def im2txt(image_path):
+    # Device configuration
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # Image preprocessing
     transform = transforms.Compose(
         [