Clear asserts, the devices are swapped in some cases. I don't think it matters.

MarkDaoust · copybara-github · commit a71f8a8b40b4 · 2022-05-18T07:17:05.000-07:00
PiperOrigin-RevId: 449475153
diff --git a/site/en/tutorials/distribute/parameter_server_training.ipynb b/site/en/tutorials/distribute/parameter_server_training.ipynb
@@ -684,8 +684,9 @@
         "assert len(emb_layer.weights) == 2\n",
         "assert emb_layer.weights[0].shape == (4, 16384)\n",
         "assert emb_layer.weights[1].shape == (4, 16384)\n",
-        "assert emb_layer.weights[0].device == \"/job:ps/replica:0/task:0/device:CPU:0\"\n",
-        "assert emb_layer.weights[1].device == \"/job:ps/replica:0/task:1/device:CPU:0\""
+        "\n",
+        "print(emb_layer.weights[0].device)\n",
+        "print(emb_layer.weights[1].device)\n"
       ]
     },
     {

Original file line number	Diff line number	Diff line change
`@@ -684,8 +684,9 @@`
`684`	`684`	`"assert len(emb_layer.weights) == 2\n",`
`685`	`685`	`"assert emb_layer.weights[0].shape == (4, 16384)\n",`
`686`	`686`	`"assert emb_layer.weights[1].shape == (4, 16384)\n",`
`687`		`- "assert emb_layer.weights[0].device == \"/job:ps/replica:0/task:0/device:CPU:0\"\n",`
`688`		`- "assert emb_layer.weights[1].device == \"/job:ps/replica:0/task:1/device:CPU:0\""`
	`687`	`+ "\n",`
	`688`	`+ "print(emb_layer.weights[0].device)\n",`
	`689`	`+ "print(emb_layer.weights[1].device)\n"`
`689`	`690`	`]`
`690`	`691`	`},`
`691`	`692`	`{`