set env variable for tensor_parallel_degree

aws · Sep 9, 2022 · 0fd8a55 · 0fd8a55
1 parent 40de44c
commit 0fd8a55
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/..._functionality/pytorch_deploy_large_GPT_model/GPT-J-6B-model-parallel-inference-DJL.ipynb b/..._functionality/pytorch_deploy_large_GPT_model/GPT-J-6B-model-parallel-inference-DJL.ipynb
@@ -299,6 +299,9 @@
     "    PrimaryContainer = {\n",
     "        'Image': fullname,\n",
     "        'ModelDataUrl': model_s3_url,\n",
+    "        'Environment': {\n",
+    "            'TENSOR_PARALLEL_DEGREE': '2'\n",
+    "        }\n",
     "    })"
    ]
   },