diff --git a/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc b/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc
index 2d7287d5c8d..f9d13938641 100644
--- a/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc
+++ b/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc
@@ -590,6 +590,10 @@ class SideEffectFinder {
     if (func_graph != nullptr) {
       return TraceTupleEffectInfo(func_graph->output(), tuple_indexes);
     }
+    // Tuple returned from a Switch call.
+    if (cnode->size() == 1 && IsPrimitiveCNode(cnode->input(0), prim::kPrimSwitch)) {
+      return TraceTupleFromSwitch(cnode->input(0)->cast<CNodePtr>(), *tuple_indexes);
+    }
     // Tuple is returned from J().
     //   %1 = J(primal)
     //   tuple = %1(args)
@@ -602,6 +606,18 @@ class SideEffectFinder {
     return {EffectInfo::kDetected, false, false, false};
   }
 
+  // Trace effect info from a Switch node that output is a tuple.
+  EffectInfo TraceTupleFromSwitch(const CNodePtr &switch_cnode, const std::stack<int64_t> &tuple_indexes) {
+    auto branches = GetSwitchBranches(switch_cnode);
+    EffectInfo info = {EffectInfo::kDetected, false, false, false};
+    for (auto &branch : branches) {
+      auto tuple_indexes_copy = tuple_indexes;
+      EffectInfo branch_info = TraceTupleEffectInfo(branch->output(), &tuple_indexes_copy);
+      info.Merge(branch_info);
+    }
+    return info;
+  }
+
   // Setup all branches according the effect info.
   void SetupEffectBranches(const EffectInfo &info, const std::vector<FuncGraphPtr> &branches) {
     // Setup monad parameters for all branches according the effect info.
diff --git a/tests/st/auto_monad/test_auto_monad_layer.py b/tests/st/auto_monad/test_auto_monad_layer.py
new file mode 100644
index 00000000000..f708725b346
--- /dev/null
+++ b/tests/st/auto_monad/test_auto_monad_layer.py
@@ -0,0 +1,81 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from tqdm import tqdm
+import numpy as np
+import mindspore as ms
+import mindspore.nn as nn
+from mindspore.dataset import NumpySlicesDataset
+from mindspore import context, Tensor
+
+context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
+
+class AutoEncoderTrainNetwork(nn.Cell):
+    def __init__(self):
+        super(AutoEncoderTrainNetwork, self).__init__()
+        self.loss_fun = nn.MSELoss()
+        self.net = nn.CellList([nn.Dense(2, 32), nn.Dense(32, 2)])
+        self.relu = nn.ReLU()
+
+    def reconstruct_sample(self, x: Tensor):
+        for _, layer in enumerate(self.net):
+            x = layer(x)
+            x = self.relu(x)
+        return x
+
+    def construct(self, x: Tensor):
+        recon_x = self.reconstruct_sample(x)
+        return self.loss_fun(recon_x, x)
+
+    def sample_2d_data(self, n_normals=2000, n_outliers=400):
+        z = np.random.randn(n_normals, 2)
+        outliers = np.random.uniform(low=-6, high=6, size=(n_outliers, 2))
+        centers = np.array([(2., 0), (-2., 0)])
+        sigma = 0.3
+        normal_points = sigma * z + centers[np.random.randint(len(centers), size=(n_normals,))]
+        return np.vstack((normal_points, outliers))
+
+    def create_synthetic_dataset(self):
+        transformed_dataset = self.sample_2d_data()
+        for dim in range(transformed_dataset.shape[1]):
+            min_val = transformed_dataset[:, dim].min()
+            max_val = transformed_dataset[:, dim].max()
+            if min_val != max_val:
+                transformed_dataset[:, dim] = (transformed_dataset[:, dim] - min_val) / (max_val - min_val)
+            elif min_val != 1:
+                transformed_dataset[:, dim] = transformed_dataset[:, dim] / min_val
+        transformed_dataset = transformed_dataset.astype(np.float32)
+        return transformed_dataset
+
+
+def test_auto_monad_layer():
+    ae_with_loss = AutoEncoderTrainNetwork()
+    transformed_dataset = ae_with_loss.create_synthetic_dataset()
+    dataloader = NumpySlicesDataset(data=(transformed_dataset,), shuffle=True)
+    dataloader = dataloader.batch(batch_size=16)
+    optim = nn.RMSProp(params=ae_with_loss.trainable_params(), learning_rate=0.002,)
+    train_net = nn.TrainOneStepCell(ae_with_loss, optim)
+    train_net.set_train()
+    gen_samples = dict()
+    num_epoch = 21
+    for epoch in tqdm(range(num_epoch)):
+        loss = []
+        for _, (batch,) in enumerate(dataloader):
+            batch = Tensor(batch, dtype=ms.float32)
+            loss_ = train_net(batch)
+            loss.append(loss_.asnumpy())
+        avg_loss = np.array(loss).mean()
+        if epoch % 10 == 0:
+            gen_samples[epoch] = ae_with_loss.reconstruct_sample(Tensor(transformed_dataset)).asnumpy()
+        print(f"epoch: {epoch}/{num_epoch}, avg loss: {avg_loss}")