    def test_create_simple_module(self):
        """
        Test simple ResNetBasicStem (without pooling layer).
        """
        for input_dim, output_dim in itertools.product((2, 3), (4, 8, 16)):
            model = ResNetBasic3DModule(
                conv=nn.Conv3d(
                    input_dim,
                    output_dim,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    bias=False,
                ),
                norm=nn.BatchNorm3d(output_dim),
                activation=nn.ReLU(),
                pool=None,
            )

            # Test forwarding.
            for tensor in TestResNetBasic3DModule._get_inputs(input_dim):
                if tensor.shape[1] != input_dim:
                    with self.assertRaises(RuntimeError):
                        output_tensor = model(tensor)
                    continue
                else:
                    output_tensor = model(tensor)

                input_shape = tensor.shape
                output_shape = output_tensor.shape
                output_shape_gt = (
                    input_shape[0],
                    output_dim,
                    input_shape[2],
                    input_shape[3],
                    input_shape[4],
                )

                self.assertEqual(
                    output_shape,
                    output_shape_gt,
                    f"Output shape {output_shape} is different from expected shape {output_shape_gt}"
                )
Test simple ResNetBasicStem (without pooling layer).
tests/unittests/models/layers/test_resnet_basic_3d_module.py