Merge pull request #72 from duggalsu/add_pytorch_arch_conditional

Add pytorch arch conditional
tattle-made · Feb 9, 2024 · 136894a · 136894a
2 parents f43646b + 2985432
commit 136894a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/src/api/Dockerfile.vid_vec_rep_resnet b/src/api/Dockerfile.vid_vec_rep_resnet
@@ -1,4 +1,4 @@
-FROM python:3.11-slim@sha256:637774748f62b832dc11e7b286e48cd716727ed04b45a0322776c01bc526afc3 AS base
+FROM python:3.11-slim-bullseye AS base
 RUN apt-get update \
     && apt-get -y upgrade \
     && apt-get install -y \

diff --git a/src/api/core/operators/vid_vec_rep_resnet_requirements.txt b/src/api/core/operators/vid_vec_rep_resnet_requirements.txt
@@ -4,7 +4,24 @@
 #
 #    pip-compile --find-links=https://download.pytorch.org/whl/torch_stable.html vid_vec_rep_resnet_requirements.in
 #
+
+# For arm64 architecture
 --find-links https://download.pytorch.org/whl/cpu
+torch==2.2.0; platform_machine=='aarch64'
+    # via
+    #   -r vid_vec_rep_resnet_requirements.in
+    #   torchvision
+torchvision==0.17.0; platform_machine=='aarch64'
+    # via -r vid_vec_rep_resnet_requirements.in
+
+# For amd64 architecture
+--find-links https://download.pytorch.org/whl/torch_stable.html
+torch==2.2.0+cpu; platform_machine=='x86_64'
+    # via
+    #   -r vid_vec_rep_resnet_requirements.in
+    #   torchvision
+torchvision==0.17.0+cpu; platform_machine=='x86_64'
+    # via -r vid_vec_rep_resnet_requirements.in
 
 certifi==2024.2.2
     # via requests
@@ -67,12 +84,6 @@ sympy==1.12
     # via torch
 textual==0.50.0
     # via memray
-torch
-    # via
-    #   -r vid_vec_rep_resnet_requirements.in
-    #   torchvision
-torchvision
-    # via -r vid_vec_rep_resnet_requirements.in
 typing-extensions==4.9.0
     # via
     #   textual