fix: remove gpu_model constraint if model is not specified

ai4os · Jan 18, 2024 · 3d11fdb · 3d11fdb
1 parent d31e3f5
commit 3d11fdb
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ai4papi/routers/v1/deployments/modules.py b/ai4papi/routers/v1/deployments/modules.py
@@ -252,11 +252,11 @@ def create_deployment(
     # Modify the GPU section
     if user_conf['hardware']['gpu_num'] <= 0:
         # Delete GPU section in CPU deployments
-        del usertask['Resources']['Devices']
+        usertask['Resources']['Devices'] = None
     else:
-        # If gpu_type not provided, remove affinity to GPU model
+        # If gpu_type not provided, remove constraint to GPU model
         if not user_conf['hardware']['gpu_type']:
-            del usertask['Resources']['Devices'][0]['Affinities']
+            usertask['Resources']['Devices'][0]['Constraints'] = None
 
     # If storage credentials not provided, remove storage-related tasks
     if not all(user_conf['storage'].values()):