code: def test_gstg_2d_input_with_mask(self) -> None:
      dim = 3
      mask = torch.tensor(
          [
              [0, 1],
              [1, 1],
              [0, 2],
          ]
      ).to(self.testing_device)
      gstg = GaussianStochasticGates(dim, mask=mask).to(
          self.testing_device
      )
  
      # shape(2,3,2)
      input_tensor = torch.tensor(
          [
              [
                  [0.0, 0.1],
                  [0.2, 0.3],
                  [0.4, 0.5],
              ],
              [
                  [0.6, 0.7],
                  [0.8, 0.9],
                  [1.0, 1.1],
              ],
          ]
      ).to(self.testing_device)
  
      gated_input, reg = gstg(input_tensor)
      expected_reg = 2.5213
  
      if self.testing_device == "cpu":
          expected_gated_input = [
              [[0.0000, 0.0198], [0.0396, 0.0594], [0.2435, 0.3708]],
              [[0.3696, 0.5954], [0.6805, 0.7655], [0.6159, 0.3921]],
          ]
      elif self.testing_device == "cuda":
          expected_gated_input = [
              [[0.0000, 0.0788], [0.1577, 0.2365], [0.0000, 0.1174]],
              [[0.0269, 0.0000], [0.0000, 0.0000], [0.0448, 0.4145]],
          ]
  
      assertTensorAlmostEqual(
          self, gated_input, expected_gated_input, mode="max"
      )
      assertTensorAlmostEqual(self, reg, expected_reg)
  
imports: import unittest
  import torch
  from captum.module.gaussian_stochastic_gates import GaussianStochasticGates
  from parameterized import parameterized_class
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
label_desc: Use image node_id 9 for calling the TestGaussianStochasticGates obj's underlying member method code with example usage: obj.test_gstg_2d_input_with_mask() without return types