diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml
index fc882f27357..fd6915277e7 100644
--- a/.github/workflows/pull.yml
+++ b/.github/workflows/pull.yml
@@ -181,7 +181,8 @@ jobs:
         { include: [
           { config: "default", shard: 1, num_shards: 2, runner: "linux.4xlarge.nvidia.gpu" },
           { config: "default", shard: 2, num_shards: 2, runner: "linux.4xlarge.nvidia.gpu" },
-          { config: "distributed", shard: 1, num_shards: 1, runner: "linux.8xlarge.nvidia.gpu" },
+          { config: "distributed", shard: 1, num_shards: 2, runner: "linux.8xlarge.nvidia.gpu" },
+          { config: "distributed", shard: 2, num_shards: 2, runner: "linux.8xlarge.nvidia.gpu" },
         ]}
 
   linux-bionic-rocm5_1-py3_7-build:
diff --git a/test/run_test.py b/test/run_test.py
index ea1a80261d0..2a1b913a10a 100644
--- a/test/run_test.py
+++ b/test/run_test.py
@@ -538,6 +538,7 @@ def test_distributed(test_module, test_directory, options):
                         backend, with_init
                     )
                 )
+            old_environ = dict(os.environ)
             os.environ["TEMP_DIR"] = tmp_dir
             os.environ["BACKEND"] = backend
             os.environ["INIT_METHOD"] = "env://"
@@ -588,6 +589,8 @@ def test_distributed(test_module, test_directory, options):
                     return return_code
             finally:
                 shutil.rmtree(tmp_dir)
+                os.environ.clear()
+                os.environ.update(old_environ)
     return 0