From 2240b5183434fc05bf45ea037aa09c14c46beb21 Mon Sep 17 00:00:00 2001
From: Rehaan Bhimani <rbhimani@microsoft.com>
Date: Wed, 7 Feb 2024 10:40:16 -0700
Subject: [PATCH] update environments for notebook (#2989)

---
 .../pytorch/distributed-training/distributed-cifar10.ipynb    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdk/python/jobs/single-step/pytorch/distributed-training/distributed-cifar10.ipynb b/sdk/python/jobs/single-step/pytorch/distributed-training/distributed-cifar10.ipynb
index aa9efef5345..949399c64db 100644
--- a/sdk/python/jobs/single-step/pytorch/distributed-training/distributed-cifar10.ipynb
+++ b/sdk/python/jobs/single-step/pytorch/distributed-training/distributed-cifar10.ipynb
@@ -254,7 +254,7 @@
     "    command=\"python read_write_data.py --input_data ${{inputs.cifar_zip}} --output_folder ${{outputs.cifar}}\",\n",
     "    inputs=inputs,\n",
     "    outputs=outputs,\n",
-    "    environment=\"AzureML-sklearn-1.0-ubuntu20.04-py38-cpu:1\",\n",
+    "    environment=\"azureml-sklearn-1.0:23\",\n",
     ")\n",
     "\n",
     "# submit the command\n",
@@ -354,7 +354,7 @@
     "    code=\"./src\",  # local path where the code is stored\n",
     "    command=\"python train.py --data-dir ${{inputs.cifar}} --epochs ${{inputs.epoch}} --batch-size ${{inputs.batchsize}} --workers ${{inputs.workers}} --learning-rate ${{inputs.lr}} --momentum ${{inputs.momen}} --print-freq ${{inputs.prtfreq}} --model-dir ${{inputs.output}}\",\n",
     "    inputs=inputs,\n",
-    "    environment=\"azureml:AzureML-pytorch-1.9-ubuntu18.04-py37-cuda11-gpu:6\",\n",
+    "    environment=\"azureml-acpt-pytorch-2.0-cuda11.7:26\",\n",
     "    instance_count=2,  # In this, only 2 node cluster was created.\n",
     "    distribution={\n",
     "        \"type\": \"PyTorch\",\n",